Skip to content

Commit 4c45400

Browse files
author
I761212
committed
TSP algo
1 parent e2a78d4 commit 4c45400

File tree

1 file changed

+65
-0
lines changed

1 file changed

+65
-0
lines changed
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
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

Comments
 (0)