Whatever-first search in Python

A generic graph traversal that keeps a bag of unvisited edges and repeatedly extracts one, with the bag's discipline giving depth-first or breadth-first search.

File: whatever_first_search.py, 37 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

whatever_first_search.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/whatever-first-search/. The authors are listed in prov/4TP999.jsonl.

from collections import deque


def whatever_first_search(
    adj: list[list[int]], s: int, bag: str = "stack"
) -> tuple[list[int], list[int | None]]:
    """Run whatever-first search from s and return (order, parent).

    The graph has the vertices 0 .. n-1 and the adjacency lists adj. The bag holds
    pairs (p, v): "v was seen from p". With bag="queue" it is a FIFO queue and the
    search is breadth-first. With bag="stack" it is a LIFO stack, and each list is
    pushed in reverse so that the first neighbour is taken first; the search is then
    depth-first. order lists the marked vertices in the order they were marked, and
    parent[v] is the p of the pair that marked v, or None. Time O(n + m).
    """
    if bag not in ("stack", "queue"):
        raise ValueError("bag must be 'stack' or 'queue'")
    n = len(adj)
    marked = [False] * n  # @line L1
    parent: list[int | None] = [None] * n
    order: list[int] = []
    items: deque[tuple[int | None, int]] = deque([(None, s)])  # @line L2
    take = items.pop if bag == "stack" else items.popleft
    while items:  # @line L3
        p, v = take()  # @line L4
        if not marked[v]:  # @line L5
            marked[v] = True  # @line L6
            parent[v] = p
            order.append(v)
            nbrs = reversed(adj[v]) if bag == "stack" else adj[v]
            for w in nbrs:  # @line L7
                if not marked[w]:  # a pair for a marked w would be skipped at L5
                    items.append((v, w))
    return order, parent  # @line L8

test_whatever_first_search.py 98 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/whatever-first-search/. The authors are listed in prov/4TP999.jsonl.

import random
from collections import deque

import pytest

from whatever_first_search import whatever_first_search


def bfs_oracle(adj: list[list[int]], s: int) -> tuple[list[int], list[int | None]]:
    """Textbook BFS: mark on enqueue, parent set at discovery."""
    parent: list[int | None] = [None] * len(adj)
    seen = [False] * len(adj)
    seen[s] = True
    order, q = [], deque([s])
    while q:
        u = q.popleft()
        order.append(u)
        for v in adj[u]:
            if not seen[v]:
                seen[v] = True
                parent[v] = u
                q.append(v)
    return order, parent


def dfs_oracle(adj: list[list[int]], s: int) -> tuple[list[int], list[int | None]]:
    """Textbook recursive DFS from s: preorder and parents."""
    parent: list[int | None] = [None] * len(adj)
    seen = [False] * len(adj)
    order: list[int] = []

    def visit(u: int) -> None:
        seen[u] = True
        order.append(u)
        for v in adj[u]:
            if not seen[v]:
                parent[v] = u
                visit(v)

    visit(s)
    return order, parent


def test_example() -> None:
    adj = [[1, 2], [0, 2], [0, 1, 3], [2]]
    assert whatever_first_search(adj, 0, "stack") == ([0, 1, 2, 3], [None, 0, 1, 2])
    assert whatever_first_search(adj, 0, "queue") == ([0, 1, 2, 3], [None, 0, 0, 2])


def test_larger_example() -> None:
    adj = [[1, 2], [0, 3], [0, 3, 4], [1, 2, 5], [2, 5], [3, 4, 6], [5], []]
    assert whatever_first_search(adj, 0, "stack") == (
        [0, 1, 3, 2, 4, 5, 6],
        [None, 0, 3, 1, 2, 4, 5, None],
    )
    assert whatever_first_search(adj, 0, "queue") == (
        [0, 1, 2, 3, 4, 5, 6],
        [None, 0, 0, 1, 2, 3, 5, None],
    )


def test_edge_cases() -> None:
    assert whatever_first_search([[]], 0) == ([0], [None])
    assert whatever_first_search([[0]], 0, "queue") == ([0], [None])
    assert whatever_first_search([[], []], 1) == ([1], [None, None])
    with pytest.raises(ValueError):
        whatever_first_search([[]], 0, "heap")


def test_long_path() -> None:
    n = 100000
    adj = [[i + 1] for i in range(n - 1)] + [[]]
    order, parent = whatever_first_search(adj, 0, "stack")
    assert order == list(range(n))
    assert parent[n - 1] == n - 2


def test_random_against_oracles() -> None:
    rng = random.Random(9137)
    for trial in range(600):
        n = rng.randint(1, 10)
        p = rng.random()
        adj: list[list[int]] = [[] for _ in range(n)]
        undirected = trial % 2 == 0
        for u in range(n):
            for v in range(n):
                if (not undirected or u < v) and rng.random() < p / 2:
                    adj[u].append(v)
                    if undirected and u != v:
                        adj[v].append(u)
        for lst in adj:
            rng.shuffle(lst)
        s = rng.randrange(n)
        assert whatever_first_search(adj, s, "queue") == bfs_oracle(adj, s)
        assert whatever_first_search(adj, s, "stack") == dfs_oracle(adj, s)

What the next level needs

the tests did not run in this build

Read the full page about Whatever-first search