|
| 1 | +""" |
| 2 | +Dynamic Programming: Travelling Salesman Problem (TSP) |
| 3 | +----------------------------------------------------- |
| 4 | +Solves the classic TSP using the Held–Karp dynamic programming approach. |
| 5 | +
|
| 6 | +Time Complexity: O(n^2 * 2^n) |
| 7 | +Space Complexity: O(n * 2^n) |
| 8 | +
|
| 9 | +Example: |
| 10 | + >>> cost = [ |
| 11 | + ... [0, 10, 15, 20], |
| 12 | + ... [10, 0, 35, 25], |
| 13 | + ... [15, 35, 0, 30], |
| 14 | + ... [20, 25, 30, 0] |
| 15 | + ... ] |
| 16 | + >>> travelling_salesman(cost) |
| 17 | + 80 |
| 18 | +""" |
| 19 | + |
| 20 | +from functools import lru_cache |
| 21 | + |
| 22 | + |
| 23 | +def travelling_salesman(cost_matrix: list[list[int]]) -> int: |
| 24 | + """ |
| 25 | + Returns the minimum travel cost for visiting all cities and returning |
| 26 | + to the starting city (0-indexed), using the Held–Karp DP approach. |
| 27 | +
|
| 28 | + Args: |
| 29 | + cost_matrix (list[list[int]]): A square matrix where cost_matrix[i][j] |
| 30 | + represents the cost of traveling from city i to city j. |
| 31 | +
|
| 32 | + Returns: |
| 33 | + int: The minimum total cost of the tour. |
| 34 | + """ |
| 35 | + n = len(cost_matrix) |
| 36 | + all_visited = (1 << n) - 1 # bitmask with all cities visited |
| 37 | + |
| 38 | + @lru_cache(maxsize=None) |
| 39 | + def dp(mask: int, pos: int) -> int: |
| 40 | + # Base case: all cities visited, return cost to go back to start |
| 41 | + if mask == all_visited: |
| 42 | + return cost_matrix[pos][0] |
| 43 | + |
| 44 | + ans = float("inf") |
| 45 | + for city in range(n): |
| 46 | + # If city not yet visited |
| 47 | + if not (mask & (1 << city)): |
| 48 | + new_cost = cost_matrix[pos][city] + dp(mask | (1 << city), city) |
| 49 | + ans = min(ans, new_cost) |
| 50 | + return ans |
| 51 | + |
| 52 | + # Start from city 0 with only it visited |
| 53 | + return dp(1, 0) |
| 54 | + |
| 55 | + |
| 56 | +if __name__ == "__main__": |
| 57 | + # Example test case |
| 58 | + cost = [ |
| 59 | + [0, 10, 15, 20], |
| 60 | + [10, 0, 35, 25], |
| 61 | + [15, 35, 0, 30], |
| 62 | + [20, 25, 30, 0], |
| 63 | + ] |
| 64 | + |
| 65 | + print("Minimum tour cost:", travelling_salesman(cost)) |
0 commit comments