Add Traveling Salesman Problem algorithms and tests

Author: MapleBauhinia

graphs/tests/test_traveling_salesman_problem.py

@@ -0,0 +1,48 @@
+import pytest
+from graphs.traveling_salesman_problem import tsp_brute_force, tsp_dp, tsp_greedy
+
+def sample_graph_1() -> list[list[int]]:
+    return [
+        [0, 29, 20],
+        [29, 0, 15],
+        [20, 15, 0],
+    ]
+
+def sample_graph_2() -> list[list[int]]:
+    return [
+        [0, 10, 15, 20],
+        [10, 0, 35, 25],
+        [15, 35, 0, 30],
+        [20, 25, 30, 0],
+    ]
+
+def test_brute_force():
+    graph = sample_graph_1()
+    assert tsp_brute_force(graph) == 64
+
+def test_dp():
+    graph = sample_graph_1()
+    assert tsp_dp(graph) == 64
+
+def test_greedy():
+    graph = sample_graph_1()
+    # The greedy algorithm does not guarantee an optimal solution; 
+    # it is necessary to verify that its output is an integer greater than 0.
+    result = tsp_greedy(graph)
+    assert isinstance(result, int)
+    assert result > 0
+
+def test_dp_larger_graph():
+    graph = sample_graph_2()
+    assert tsp_dp(graph) == 80  
+
+def test_brute_force_larger_graph():
+    graph = sample_graph_2()
+    assert tsp_brute_force(graph) == 80
+
+def test_greedy_larger_graph():
+    graph = sample_graph_2()
+    # An approximate solution cannot be represented by '==' and can only ensure that the result is reasonable.
+    result = tsp_greedy(graph)
+    assert isinstance(result, int)
+    assert result >= 80

graphs/traveling_salesman_problem.py

@@ -0,0 +1,151 @@
+from itertools import permutations
+
+def tsp_brute_force(graph: list[list[int]]) -> int:
+    """
+    Solves TSP using brute-force permutations.
+
+    Args:
+        graph: 2D list representing distances between cities.
+
+    Returns:
+        The minimal total travel distance visiting all cities exactly once and returning to the start.
+
+    Example:
+        >>> tsp_brute_force([[0, 29, 20], [29, 0, 15], [20, 15, 0]])
+        64
+    """
+    n = len(graph)
+    # Apart from other cities aside from City 0, City 0 serves as the starting point.
+    nodes = list(range(1, n)) 
+    min_path = float('inf')
+    
+    # Enumerate all the permutations from city 1 to city n-1.
+    for perm in permutations(nodes):
+        # Construct a complete path: 
+        # starting from point 0, visit in the order of arrangement, and then return to point 0.
+        path = [0] + list(perm) + [0]
+
+        # Calculate the total distance of the path. 
+        # Update the shortest path.
+        total_cost = sum(graph[path[i]][path[i + 1]] for i in range(n))
+        min_path = min(min_path, total_cost)
+
+    return min_path
+
+def tsp_dp(graph: list[list[int]]) -> int:
+    """
+    Solves the Traveling Salesman Problem using Held-Karp dynamic programming.
+
+    Args:
+        graph: A 2D list representing distances between cities (n x n matrix).
+
+    Returns:
+        The minimum cost to visit all cities exactly once and return to the origin.
+
+    Example:
+        >>> tsp_dp([[0, 29, 20], [29, 0, 15], [20, 15, 0]])
+        64
+    """
+    n = len(graph)
+    # Create a dynamic programming table of size (2^n) x n.
+    # Noting: 1 << n  = 2^n
+    # dp[mask][i] represents the shortest path starting from city 0, passing through the cities in the mask, and ultimately ending at city i.
+    dp = [[float('inf')] * n for _ in range(1 << n)]
+    # Initial state: only city 0 is visited, and the path length is 0.
+    dp[1][0] = 0  
+
+    for mask in range(1 << n):
+        # The mask indicates which cities have been visited.
+        for u in range(n):
+            if not (mask & (1 << u)):
+                # If the city u is not included in the mask, skip it.
+                continue
+
+            for v in range(n):
+                # City v has not been accessed and is different from city u.
+                if mask & (1 << v) or u == v:
+                    continue
+
+                # New State: Transition to city v
+                # State Transition: From city u to city v, updating the shortest path.
+                next_mask = mask | (1 << v)
+                dp[next_mask][v] = min(dp[next_mask][v], dp[mask][u] + graph[u][v])
+    
+    # After completing visits to all cities, return to city 0 and obtain the minimum value.
+    return min(dp[(1 << n) - 1][i] + graph[i][0] for i in range(1, n))
+
+def tsp_greedy(graph: list[list[int]]) -> int:
+    """
+    Solves TSP approximately using the nearest neighbor heuristic.
+    Warming: This algorithm is not guaranteed to find the optimal solution! But it is fast and applicable to any input size.
+
+    Args:
+        graph: 2D list representing distances between cities.
+
+    Returns:
+        The total distance of the approximated TSP route.
+
+    Example:
+        >>> tsp_greedy([[0, 29, 20], [29, 0, 15], [20, 15, 0]])
+        64
+    """
+    n = len(graph)
+    visited = [False] * n  # Mark whether each city has been visited.
+    path = [0]
+    total_cost = 0
+    visited[0] = True  # Start from city 0.
+    current = 0  # Current city.
+
+    for _ in range(n - 1):
+        # Find the nearest city to the current location that has not been visited.
+        next_city = min(
+            ((city, cost) for city, cost in enumerate(graph[current]) if not visited[city] and city != current),
+            key=lambda x: x[1],
+            default=(None, float('inf'))
+        )[0]
+
+        # If no such city exists, break the loop.
+        if next_city is None:
+            break
+
+        # Update the total cost and the current city.
+        # Mark the city as visited.
+        # Append the city to the path.
+        total_cost += graph[current][next_city]
+        visited[next_city] = True
+        current = next_city
+        path.append(current)
+
+    # Back to start
+    total_cost += graph[current][0]
+    path.append(0)
+
+    return total_cost
+
+
+def test_tsp_example():
+    graph = [[0, 29, 20], [29, 0, 15], [20, 15, 0]]
+
+    result = tsp_brute_force(graph)
+    if result != 64:
+        raise Exception('tsp_brute_force Incorrect result')
+    else:
+        print('Test passed')
+    
+    result = tsp_dp(graph)
+    if result != 64:
+        raise Exception('tsp_dp Incorrect result')
+    else:
+        print("Test passed")
+    
+    result = tsp_greedy(graph)
+    if result != 64:
+        if result < 0:
+            raise Exception('tsp_greedy Incorrect result')
+        else:
+            print("tsp_greedy gets an approximate result.")
+    else:
+        print('Test passed')
+
+if __name__ == '__main__':
+    test_tsp_example()