Graphs | Solving the Travelling Salesperson Problem (TSP) in Python
From Brute Force to Held-Karp: Exploring Three Algorithms to Conquer TSP with Python
The Travelling Salesperson Problem (TSP) involves finding the shortest route that visits a set of places and returns to the origin.
In TSP, a “sales person” starts at a home city, visits a list of given cities exactly once, and returns to the home city, minimizing the total travel distance or cost. This problem can be applied to various practical scenarios, including logistics, planning, and manufacturing.
Despite its straightforward description, TSP is a cornerstone of combinatorial optimization and an NP-hard problem: solving it efficiently for large numbers of cities remains a computational challenge. At its heart, TSP is deeply tied to graph theory, a branch of mathematics that studies the relationships between objects represented as nodes (or vertices) and edges.
In this post, we’ll use Python to implement three solutions — Brute Force, Nearest Neighbor, and Held-Karp. We’ll generate cities programmatically, measure performance, and visualize the results while letting you tweak the number of cities to see how each algorithm scales.
The Code: A Walkthrough
Here’s the full implementation, followed by a breakdown of each section:
import itertools
import networkx as nx
import matplotlib.pyplot as plt
from itertools import combinations
import time
import random
import numpy as np
# Function to create a graph with a specified number of nodes
def create_graph(num_nodes, seed=42):
random.seed(seed)
G = nx.Graph()
cities = {i: (random.uniform(0, 10), random.uniform(0, 10)) for i in range(num_nodes)}
for city, pos in cities.items():
G.add_node(city, pos=pos)
for i in range(num_nodes):
for j in range(i + 1, num_nodes):
x1, y1 = cities[i]
x2, y2 = cities[j]
distance = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
G.add_edge(i, j, weight=distance)
return G, cities
def route_distance(G, route):
total = 0
for i in range(len(route) - 1):
total += G[route[i]][route[i + 1]]['weight']
total += G[route[-1]][route[0]]['weight']
return total
def solve_tsp_brute_force(G):
num_nodes = G.number_of_nodes()
other_cities = list(range(1, num_nodes))
all_routes = itertools.permutations(other_cities)
best_route = None
min_distance = float('inf')
for perm in all_routes:
current_route = [0] + list(perm)
current_distance = route_distance(G, current_route)
if current_distance < min_distance:
min_distance = current_distance
best_route = current_route
return best_route, min_distance
def solve_tsp_nearest_neighbor(G):
num_nodes = G.number_of_nodes()
route = [0]
unvisited = set(range(1, num_nodes))
while unvisited:
current_city = route[-1]
next_city = min(unvisited, key=lambda x: G[current_city][x]['weight'])
route.append(next_city)
unvisited.remove(next_city)
return route, route_distance(G, route)
def solve_tsp_held_karp(G):
n = G.number_of_nodes()
dp = {}
for i in range(1, n):
dp[(1 << i, i)] = (G[0][i]['weight'], 0)
for size in range(2, n):
for subset in combinations(range(1, n), size):
bits = sum(1 << i for i in subset)
for end in subset:
prev_bits = bits & ~(1 << end)
min_cost = float('inf')
min_prev = None
for prev in subset:
if prev == end or (prev_bits, prev) not in dp:
continue
cost = dp[(prev_bits, prev)][0] + G[prev][end]['weight']
if cost < min_cost:
min_cost = cost
min_prev = prev
dp[(bits, end)] = (min_cost, min_prev)
all_bits = (1 << n) - 2
min_cost = float('inf')
final_end = None
for end in range(1, n):
if (all_bits, end) in dp:
cost = dp[(all_bits, end)][0] + G[end][0]['weight']
if cost < min_cost:
min_cost = cost
final_end = end
route = [final_end]
current_bits = all_bits
while current_bits:
last = route[-1]
cost, prev = dp[(current_bits, last)]
route.append(prev)
current_bits &= ~(1 << last)
route = route[::-1]
route[0] = 0
return route, min_cost
def plot_tsp(G, route, distance, title, time_taken):
pos = nx.get_node_attributes(G, 'pos')
plt.figure(figsize=(10, 8))
nx.draw_networkx_nodes(G, pos, node_color='lightblue', node_size=500)
nx.draw_networkx_edges(G, pos, edge_color='gray', alpha=0.3)
nx.draw_networkx_labels(G, pos, font_size=12, font_weight='bold')
best_edges = [(route[i], route[i + 1]) for i in range(len(route) - 1)]
best_edges.append((route[-1], route[0]))
nx.draw_networkx_edges(G, pos, edgelist=best_edges, edge_color='red', width=2)
edge_labels = {(i, j): f"{G[i][j]['weight']:.2f}" for i, j in G.edges()}
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=10)
plt.title(f"{title}\nRoute: {route}, Distance: {distance:.2f}, Time: {time_taken:.6f}s", fontsize=14)
plt.axis('off')
plt.show()
def run_tsp_analysis(num_nodes):
G, cities = create_graph(num_nodes)
print(f"\nAnalyzing TSP with {num_nodes} nodes:")
start_time = time.time()
bf_route, bf_distance = solve_tsp_brute_force(G)
bf_time = time.time() - start_time
print(f"Brute Force - Route: {bf_route}, Distance: {bf_distance:.2f}, Time: {bf_time:.6f}s")
start_time = time.time()
nn_route, nn_distance = solve_tsp_nearest_neighbor(G)
nn_time = time.time() - start_time
print(f"Nearest Neighbor - Route: {nn_route}, Distance: {nn_distance:.2f}, Time: {nn_time:.6f}s")
start_time = time.time()
hk_route, hk_distance = solve_tsp_held_karp(G)
hk_time = time.time() - start_time
print(f"Held-Karp - Route: {hk_route}, Distance: {hk_distance:.2f}, Time: {hk_time:.6f}s")
plot_tsp(G, bf_route, bf_distance, f"TSP Solution (Brute Force, {num_nodes} nodes)", bf_time)
plot_tsp(G, nn_route, nn_distance, f"TSP Solution (Nearest Neighbor, {num_nodes} nodes)", nn_time)
plot_tsp(G, hk_route, hk_distance, f"TSP Solution (Held-Karp, {num_nodes} nodes)", hk_time)
run_tsp_analysis(10)
Imports
We start with the necessary libraries:
itertools
: For generating permutations (Brute Force) and combinations (Held-Karp).networkx
: To build and manipulate our graph.matplotlib.pyplot
: For visualization.time
: To measure execution time.random
andnumpy
: For generating random coordinates and calculating distances.
These tools give us everything we need to model TSP as a graph problem and analyze it.
Creating the Graph Programmatically
The create_graph(num_nodes, seed=42)
function is where the magic begins:
- Nodes: We generate
num_nodes
cities with random (x, y) coordinates in a 10x10 grid. Theseed
ensures reproducibility. - Edges: We create a complete graph (every city connects to every other) and compute edge weights as Euclidean distances with
np.sqrt((x2-x1)**2 + (y2-y1)**2)
. - Return: It outputs the
networkx
graph G and the cities dictionary, though we only use G later.
This flexibility lets you test TSP with any number of cities — just change num_nodes
.
Calculating Route Distance
The route_distance(G,route)
function computes the total distance of a route:
- Iterates over consecutive pairs in the route, summing edge weights from G.
- Adds the return trip from the last city to the start to complete the cycle.
This is shared across all algorithms, ensuring consistent distance calculations.
Algorithm 1: Brute Force
solve_tsp_brute_force(G)
takes the simplest approach:
- Generates all permutations of cities (excluding the fixed start at 0) using
itertools.permutations
. - Computes the distance for each route and keeps the shortest.
- Complexity: 𝑂(𝑛!), checking every possible tour. It’s guaranteed to find the optimal solution but becomes impractical as 𝑛 grows.
Algorithm 2: Nearest Neighbor
solve_tsp_nearest_neighbor(G)
uses a greedy heuristic:
- Starts at city 0 and repeatedly picks the closest unvisited city until all are visited.
- Uses a
set
for efficient tracking of unvisited nodes. - Complexity: 𝑂(𝑛2), fast and simple.
It’s quick but approximate — sometimes it stumbles into the optimum, sometimes it doesn’t.
Algorithm 3: Held-Karp
solve_tsp_nearest_held_karp(G)
is the dynamic programming star:
- Builds a table
dp[(S,i)]
whereS
is a bitmask of visited cities (excluding 0), andi
is the ending city. - Base case: Distances from 0 to each city.
- Recurrence: For each subset size, computes the minimum cost to end at each city in the subset.
- Reconstructs the route by backtracking through parent pointers.
- Complexity: O(n^2 . 2^n), , exponential but far better than O(n!).
It’s exact and more efficient than Brute Force for moderate n.
Visualization
plot_tsp(G, route, distance, title, time_taken)
brings the solutions to life:
- Plots nodes at their (x, y) coordinates with
solve_tsp_networkx
. - Draws all edges in faint gray and the solution route in bold red.
- Labels edges with weights and includes the route, distance, and time in the title.
Each algorithm gets its own plot, making comparisons visual and intuitive.
Running the Analysis
run_tsp_analysis(num_nodes)
ties it all together:
- Creates the graph with the specified number of nodes.
- Runs each algorithm, timing them with
time.time()
. - Prints the route, distance, and execution time.
- Visualizes all three solutions.
You can call it with any n, like run_tsp_analysis(8)
or run_tsp_analysis(12)
.
Final Thoughts
Each algorithm tells a story of trade-offs: raw power versus speed versus ingenuity. All three algorithms will find the same optimal route for small node numbers, but for larger n, Brute Force slows dramatically, and Nearest Neighbor may diverge from the optimum.
Want to push this further? Add more nodes, experiment with non-Euclidean weights, or combine Nearest Neighbor with 2-Opt for better approximations. TSP is a rabbit hole of optimization.