Held-Karp algorithm in Python
Solve the exact travelling salesman problem in O(n^2 2^n) time by DP over subsets.
- File
held_karp_algorithm.py, 53 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/held-karp-algorithm/. The authors are listed in prov/X3NDEJ.jsonl.
from typing import Sequence
INF = float("inf")
def held_karp(d: Sequence[Sequence[float]]) -> tuple[float, list[int]]:
"""Return (length, tour) of a shortest travelling salesman tour.
d is an n by n matrix of distances, n >= 1; d[i][j] is the cost of the
edge from city i to city j and need not be symmetric. The tour starts at
city 0, lists every city once, and the edge back to 0 is implied.
Time Theta(n^2 2^n), space Theta(n 2^n).
"""
n = len(d)
if n == 1: # @line L1
return 0, [0]
m = n - 1 # cities 1 .. n-1 use bits 0 .. n-2
full = (1 << m) - 1
c = [[INF] * n for _ in range(1 << m)]
p = [[-1] * n for _ in range(1 << m)]
for j in range(1, n): # @line L2
c[1 << (j - 1)][j] = d[0][j]
p[1 << (j - 1)][j] = 0
for s in range(1, full + 1): # @line L3
if s & (s - 1) == 0:
continue # sets of size 1 were filled by line L2
for j in range(1, n): # @line L4
bit = 1 << (j - 1)
if not s & bit:
continue
r = s ^ bit
best_val, best_i = INF, -1
for i in range(1, n): # @line L5
if r & (1 << (i - 1)):
v = c[r][i] + d[i][j]
if v < best_val:
best_val, best_i = v, i
c[s][j] = best_val
p[s][j] = best_i # @line L6
best, last = INF, -1
for j in range(1, n): # @line L7
v = c[full][j] + d[j][0]
if v < best:
best, last = v, j
s, j, tour = full, last, [] # @line L8
while j != 0: # @line L9
tour.append(j) # @line L10
s, j = s ^ (1 << (j - 1)), p[s][j] # @line L11
tour.reverse()
return best, [0] + tour # @line L12
test_held_karp_algorithm.py 54 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/held-karp-algorithm/. The authors are listed in prov/X3NDEJ.jsonl.
import itertools
import random
from held_karp_algorithm import held_karp
def tour_length(d, tour):
return sum(d[tour[k]][tour[(k + 1) % len(tour)]] for k in range(len(tour)))
def brute(d):
n = len(d)
return min(tour_length(d, [0] + list(p)) for p in itertools.permutations(range(1, n)))
def check(d):
length, tour = held_karp(d)
assert sorted(tour) == list(range(len(d)))
assert tour[0] == 0
assert tour_length(d, tour) == length
assert length == brute(d)
def test_example() -> None:
d = [[0, 2, 9, 10], [1, 0, 6, 4], [15, 7, 0, 8], [6, 3, 12, 0]]
assert held_karp(d) == (21, [0, 2, 3, 1])
def test_small_sizes() -> None:
assert held_karp([[0]]) == (0, [0])
assert held_karp([[0, 3], [5, 0]]) == (8, [0, 1])
check([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
def test_equal_and_negative_weights() -> None:
check([[0] * 5 for _ in range(5)])
check([[0, -3, 4, 1], [2, 0, -1, 5], [-2, 7, 0, 0], [3, -4, 6, 0]])
def test_large_values() -> None:
big = 1 << 60
d = [[0, big, 1, big], [1, 0, big, big], [big, big, 0, 1], [big, 1, big, 0]]
assert held_karp(d) == (4, [0, 2, 3, 1])
def test_against_brute_force() -> None:
rng = random.Random(20240601)
for _ in range(600):
n = rng.randint(1, 7)
d = [[0 if i == j else rng.randint(-5, 30) for j in range(n)] for i in range(n)]
check(d)
What the next level needs
- the tests did not run in this build