Skip to content

Commit c32a022

Browse files
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent e39c3ce commit c32a022

File tree

2 files changed

+37
-21
lines changed

2 files changed

+37
-21
lines changed

graphs/tests/test_traveling_salesman_problem.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
import pytest
22
from graphs.traveling_salesman_problem import tsp_brute_force, tsp_dp, tsp_greedy
33

4+
45
def sample_graph_1() -> list[list[int]]:
56
return [
67
[0, 29, 20],
78
[29, 0, 15],
89
[20, 15, 0],
910
]
1011

12+
1113
def sample_graph_2() -> list[list[int]]:
1214
return [
1315
[0, 10, 15, 20],
@@ -16,31 +18,37 @@ def sample_graph_2() -> list[list[int]]:
1618
[20, 25, 30, 0],
1719
]
1820

21+
1922
def test_brute_force():
2023
graph = sample_graph_1()
2124
assert tsp_brute_force(graph) == 64
2225

26+
2327
def test_dp():
2428
graph = sample_graph_1()
2529
assert tsp_dp(graph) == 64
2630

31+
2732
def test_greedy():
2833
graph = sample_graph_1()
29-
# The greedy algorithm does not guarantee an optimal solution;
34+
# The greedy algorithm does not guarantee an optimal solution;
3035
# it is necessary to verify that its output is an integer greater than 0.
3136
# An approximate solution cannot be represented by '==' and can only ensure that the result is reasonable.
3237
result = tsp_greedy(graph)
3338
assert isinstance(result, int)
3439
assert result >= 64
3540

41+
3642
def test_dp_larger_graph():
3743
graph = sample_graph_2()
38-
assert tsp_dp(graph) == 80
44+
assert tsp_dp(graph) == 80
45+
3946

4047
def test_brute_force_larger_graph():
4148
graph = sample_graph_2()
4249
assert tsp_brute_force(graph) == 80
4350

51+
4452
def test_greedy_larger_graph():
4553
graph = sample_graph_2()
4654
result = tsp_greedy(graph)

graphs/traveling_salesman_problem.py

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from itertools import permutations
22

3+
34
def tsp_brute_force(graph: list[list[int]]) -> int:
45
"""
56
Solves TSP using brute-force permutations.
@@ -16,22 +17,23 @@ def tsp_brute_force(graph: list[list[int]]) -> int:
1617
"""
1718
n = len(graph)
1819
# Apart from other cities aside from City 0, City 0 serves as the starting point.
19-
nodes = list(range(1, n))
20-
min_path = float('inf')
21-
20+
nodes = list(range(1, n))
21+
min_path = float("inf")
22+
2223
# Enumerate all the permutations from city 1 to city n-1.
2324
for perm in permutations(nodes):
24-
# Construct a complete path:
25+
# Construct a complete path:
2526
# starting from point 0, visit in the order of arrangement, and then return to point 0.
2627
path = [0] + list(perm) + [0]
2728

28-
# Calculate the total distance of the path.
29+
# Calculate the total distance of the path.
2930
# Update the shortest path.
3031
total_cost = sum(graph[path[i]][path[i + 1]] for i in range(n))
3132
min_path = min(min_path, total_cost)
3233

3334
return min_path
3435

36+
3537
def tsp_dp(graph: list[list[int]]) -> int:
3638
"""
3739
Solves the Traveling Salesman Problem using Held-Karp dynamic programming.
@@ -50,9 +52,9 @@ def tsp_dp(graph: list[list[int]]) -> int:
5052
# Create a dynamic programming table of size (2^n) x n.
5153
# Noting: 1 << n = 2^n
5254
# dp[mask][i] represents the shortest path starting from city 0, passing through the cities in the mask, and ultimately ending at city i.
53-
dp = [[float('inf')] * n for _ in range(1 << n)]
55+
dp = [[float("inf")] * n for _ in range(1 << n)]
5456
# Initial state: only city 0 is visited, and the path length is 0.
55-
dp[1][0] = 0
57+
dp[1][0] = 0
5658

5759
for mask in range(1 << n):
5860
# The mask indicates which cities have been visited.
@@ -70,10 +72,11 @@ def tsp_dp(graph: list[list[int]]) -> int:
7072
# State Transition: From city u to city v, updating the shortest path.
7173
next_mask = mask | (1 << v)
7274
dp[next_mask][v] = min(dp[next_mask][v], dp[mask][u] + graph[u][v])
73-
75+
7476
# After completing visits to all cities, return to city 0 and obtain the minimum value.
7577
return min(dp[(1 << n) - 1][i] + graph[i][0] for i in range(1, n))
7678

79+
7780
def tsp_greedy(graph: list[list[int]]) -> int:
7881
"""
7982
Solves TSP approximately using the nearest neighbor heuristic.
@@ -101,9 +104,13 @@ def tsp_greedy(graph: list[list[int]]) -> int:
101104
for _ in range(n - 1):
102105
# Find the nearest city to the current location that has not been visited.
103106
next_city = min(
104-
((city, cost) for city, cost in enumerate(graph[current]) if not visited[city] and city != current),
107+
(
108+
(city, cost)
109+
for city, cost in enumerate(graph[current])
110+
if not visited[city] and city != current
111+
),
105112
key=lambda x: x[1],
106-
default=(None, float('inf'))
113+
default=(None, float("inf")),
107114
)[0]
108115

109116
# If no such city exists, break the loop.
@@ -135,24 +142,25 @@ def test_tsp_example():
135142

136143
result = tsp_brute_force(graph)
137144
if result != 80:
138-
raise Exception('tsp_brute_force Incorrect result')
145+
raise Exception("tsp_brute_force Incorrect result")
139146
else:
140-
print('Test passed')
141-
147+
print("Test passed")
148+
142149
result = tsp_dp(graph)
143150
if result != 80:
144-
raise Exception('tsp_dp Incorrect result')
151+
raise Exception("tsp_dp Incorrect result")
145152
else:
146153
print("Test passed")
147-
154+
148155
result = tsp_greedy(graph)
149156
if result != 80:
150157
if result < 0:
151-
raise Exception('tsp_greedy Incorrect result')
158+
raise Exception("tsp_greedy Incorrect result")
152159
else:
153160
print("tsp_greedy gets an approximate result.")
154161
else:
155-
print('Test passed')
162+
print("Test passed")
163+
156164

157-
if __name__ == '__main__':
158-
test_tsp_example()
165+
if __name__ == "__main__":
166+
test_tsp_example()

0 commit comments

Comments
 (0)