Uniform-cost search in Python
Expands the frontier state with the smallest path cost so far; optimal for non-negative step costs.
- File
uniform_cost_search.py, 60 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/uniform-cost-search/. The authors are listed in prov/81NJ23.jsonl.
import heapq
import itertools
from collections.abc import Callable, Hashable, Iterable
from typing import NamedTuple, Optional, TypeVar
S = TypeVar("S", bound=Hashable)
class Node(NamedTuple):
state: Hashable
parent: Optional["Node"]
g: float
def _path(n: Node) -> list:
out = []
m: Optional[Node] = n
while m is not None:
out.append(m.state)
m = m.parent
out.reverse()
return out
def uniform_cost_search(
start: S,
successors: Callable[[S], Iterable[tuple[S, float]]],
is_goal: Callable[[S], bool],
) -> Optional[tuple[list[S], float]]:
"""Return (path, cost) of a cheapest path from start to a goal, or None.
successors(s) yields pairs (next_state, cost) with cost >= 0. States must be
hashable. The goal test is done when a node leaves the frontier, so the cost
is optimal. The search stops on every finite space, and on an infinite space
with a finite branching factor when the costs are at least some eps > 0 and
a goal is reachable. Cost: O((V + E) log V) on an explicit graph with V
reachable states and E transitions; O(N log N) for N generated nodes.
"""
tie = itertools.count() # breaks ties in the heap by insertion order
node = Node(start, None, 0) # @line L1
frontier: list[tuple[float, int, Node]] = [(0, next(tie), node)] # @line L2
reached: dict[Hashable, Node] = {start: node} # @line L3
while frontier: # @line L4
_, _, n = heapq.heappop(frontier) # @line L5
if n.g > reached[n.state].g: # @line L6
continue
if is_goal(n.state): # @line L7
return _path(n), n.g
for s2, c in successors(n.state): # @line L8
if c < 0:
raise ValueError("uniform-cost search needs non-negative action costs")
g2 = n.g + c # @line L9
old = reached.get(s2)
if old is None or g2 < old.g: # @line L10
reached[s2] = Node(s2, n, g2) # @line L11
heapq.heappush(frontier, (g2, next(tie), reached[s2])) # @line L12
return None # @line L13
test_uniform_cost_search.py 103 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/uniform-cost-search/. The authors are listed in prov/81NJ23.jsonl.
import itertools
import random
import pytest
from uniform_cost_search import uniform_cost_search
EXAMPLE = {
"S": [("A", 1), ("G", 12)],
"A": [("B", 4), ("C", 1)],
"B": [("D", 3)],
"C": [("D", 1), ("G", 2)],
"D": [("G", 3)],
"G": [],
}
def run(graph, start, goals):
return uniform_cost_search(start, lambda s: graph.get(s, []), lambda s: s in goals)
def path_cost(graph, path):
total = 0
for u, v in zip(path, path[1:]):
total += min(c for w, c in graph[u] if w == v)
return total
def oracle(graph, start, goals):
"""Least cost to a goal by Bellman-Ford relaxation, or None."""
nodes = set(graph) | {v for es in graph.values() for v, _ in es} | {start}
dist = {v: float("inf") for v in nodes}
dist[start] = 0
for _ in range(len(nodes)):
for u, es in graph.items():
for v, c in es:
if dist[u] + c < dist[v]:
dist[v] = dist[u] + c
best = min((dist[g] for g in goals if g in dist), default=float("inf"))
return None if best == float("inf") else best
def test_example() -> None:
assert run(EXAMPLE, "S", {"G"}) == (["S", "A", "C", "G"], 4)
def test_start_is_goal() -> None:
assert run(EXAMPLE, "S", {"S"}) == (["S"], 0)
def test_unreachable() -> None:
assert run(EXAMPLE, "G", {"S"}) is None
assert run({"X": []}, "X", set()) is None
def test_cycles_and_zero_costs() -> None:
g = {"a": [("b", 0)], "b": [("a", 0), ("c", 0)], "c": [("a", 0)]}
assert run(g, "a", {"z"}) is None
g["c"].append(("z", 5))
assert run(g, "a", {"z"}) == (["a", "b", "c", "z"], 5)
def test_negative_cost_rejected() -> None:
with pytest.raises(ValueError):
run({"a": [("b", -1)]}, "a", {"b"})
def test_implicit_infinite_space() -> None:
# states are integers; n -> n + 1 costs 1, n -> 2n costs 1; a cheapest way from 1 to 37
res = uniform_cost_search(1, lambda n: [(n + 1, 1), (2 * n, 1)], lambda n: n == 37)
assert res is not None
path, cost = res
assert cost == 7 and path[0] == 1 and path[-1] == 37 and len(path) == 8
def test_float_costs() -> None:
g = {"s": [("a", 0.1), ("t", 0.35)], "a": [("b", 0.1)], "b": [("t", 0.1)]}
path, cost = run(g, "s", {"t"})
assert path == ["s", "a", "b", "t"]
assert abs(cost - 0.3) < 1e-12
def test_against_bellman_ford() -> None:
rng = random.Random(81)
for _ in range(800):
n = rng.randrange(1, 9)
graph = {v: [] for v in range(n)}
for u, v in itertools.product(range(n), repeat=2):
if rng.random() < 0.3:
graph[u].append((v, rng.randrange(0, 6)))
goals = {v for v in range(n) if rng.random() < 0.25}
want = oracle(graph, 0, goals)
got = run(graph, 0, goals)
if want is None:
assert got is None
else:
path, cost = got
assert cost == want
assert path[0] == 0 and path[-1] in goals
assert path_cost(graph, path) == cost
What the next level needs
- the tests did not run in this build