0-1 BFS in Python
Shortest paths in graphs with weights 0 or 1 using a deque, in linear time.
- File
zero_one_bfs.py, 42 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/zero-one-bfs/. The authors are listed in prov/T8G377.jsonl.
from collections import deque
def zero_one_bfs(adj: list[list[tuple[int, int]]], s: int) -> tuple[list[int | None], list[int | None]]:
"""Single-source shortest paths when every arc weight is 0 or 1.
adj[u] lists the arcs (v, w) out of u, for vertices 0 .. n-1, with w in {0, 1}.
Returns (dist, parent): dist[v] is the distance from s, or None when v is not
reachable; parent[v] is the vertex before v on one shortest path (None for s and
for unreachable vertices). Raises ValueError on a weight other than 0 or 1, so a
negative arc (and so a negative cycle) is rejected. Time O(n + m), extra space O(n + m).
"""
n = len(adj)
for u in range(n):
for _, w in adj[u]:
if w not in (0, 1):
raise ValueError(f"arc weight {w} is not 0 or 1")
dist: list[int | None] = [None] * n # @line L1: None stands for infinity
parent: list[int | None] = [None] * n
done = [False] * n
dist[s] = 0 # @line L2
dq: deque[int] = deque([s])
while dq: # @line L3
u = dq.popleft() # @line L4
if done[u]: # @line L5: a stale copy of u
continue
done[u] = True # @line L6
du = dist[u]
assert du is not None
for v, w in adj[u]: # @line L7
dv = dist[v]
if dv is None or du + w < dv: # @line L8
dist[v] = du + w # @line L9
parent[v] = u
if w == 0: # @line L10
dq.appendleft(v)
else:
dq.append(v) # @line L11
return dist, parent # @line L12
test_zero_one_bfs.py 89 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/zero-one-bfs/. The authors are listed in prov/T8G377.jsonl.
import heapq
import random
import pytest
from zero_one_bfs import zero_one_bfs
S, A, B, C, D, T = range(6)
EXAMPLE = [(S, A, 1), (S, B, 0), (B, A, 0), (A, C, 1), (B, D, 1), (C, T, 0), (D, C, 0), (D, T, 1)]
def build(n: int, arcs: list[tuple[int, int, int]]) -> list[list[tuple[int, int]]]:
adj: list[list[tuple[int, int]]] = [[] for _ in range(n)]
for u, v, w in arcs:
adj[u].append((v, w))
return adj
def dijkstra(adj: list[list[tuple[int, int]]], s: int) -> list[int | None]:
dist: list[int | None] = [None] * len(adj)
pq = [(0, s)]
while pq:
d, u = heapq.heappop(pq)
if dist[u] is not None:
continue
dist[u] = d
for v, w in adj[u]:
if dist[v] is None:
heapq.heappush(pq, (d + w, v))
return dist
def check_tree(adj: list[list[tuple[int, int]]], s: int, dist: list[int | None], parent: list[int | None]) -> None:
for v in range(len(adj)):
if v == s or dist[v] is None:
assert parent[v] is None
continue
u = parent[v]
assert u is not None and dist[u] is not None
assert any(x == v and dist[u] + w == dist[v] for x, w in adj[u])
def test_example() -> None:
adj = build(6, EXAMPLE)
dist, parent = zero_one_bfs(adj, S)
assert dist == [0, 0, 0, 1, 1, 1]
assert parent == [None, B, S, A, B, C]
check_tree(adj, S, dist, parent)
def test_unreachable() -> None:
adj = build(4, [(0, 1, 1), (2, 3, 0), (3, 0, 1)])
dist, parent = zero_one_bfs(adj, 0)
assert dist == [0, 1, None, None]
assert parent == [None, 0, None, None]
def test_single_vertex_and_self_loop() -> None:
assert zero_one_bfs([[]], 0) == ([0], [None])
assert zero_one_bfs([[(0, 0), (0, 1)]], 0) == ([0], [None])
def test_zero_cycle() -> None:
adj = build(3, [(0, 1, 1), (1, 2, 0), (2, 1, 0)])
dist, parent = zero_one_bfs(adj, 0)
assert dist == [0, 1, 1]
assert parent == [None, 0, 1]
def test_rejects_negative_and_large_weights() -> None:
with pytest.raises(ValueError):
zero_one_bfs(build(2, [(0, 1, 0), (1, 0, -1)]), 0) # a negative cycle
with pytest.raises(ValueError):
zero_one_bfs(build(2, [(0, 1, 2)]), 0)
def test_random_against_dijkstra() -> None:
rng = random.Random(1201)
for _ in range(600):
n = rng.randint(1, 9)
arcs = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 1)) for _ in range(rng.randint(0, 3 * n))]
adj = build(n, arcs)
s = rng.randrange(n)
dist, parent = zero_one_bfs(adj, s)
assert dist == dijkstra(adj, s)
check_tree(adj, s, dist, parent)
What the next level needs
- the tests did not run in this build