BK-tree in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bk-tree/. The authors are listed in prov/WM7ED9.jsonl.

from collections.abc import Callable, Hashable, Iterable
from typing import Generic, TypeVar

T = TypeVar("T", bound=Hashable)


def levenshtein(a: str, b: str) -> int:
    """Edit distance with unit-cost insertion, deletion and substitution. O(|a| |b|) time."""
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        cur = [i]
        for j, cb in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != cb)))
        prev = cur
    return prev[-1]


class _Node(Generic[T]):
    __slots__ = ("item", "children")

    def __init__(self, item: T) -> None:
        self.item = item
        self.children: dict[int, _Node[T]] = {}


class BKTree(Generic[T]):
    """Burkhard-Keller tree for a metric with nonnegative integer values.

    Invariant: every item in the subtree of t.children[i] is at distance exactly i from t.item.
    Insert: O(h) distance evaluations. Range and nearest: O(n) evaluations in the worst case.
    """

    def __init__(self, dist: Callable[[T, T], int] = levenshtein, items: Iterable[T] = ()) -> None:
        self.dist = dist
        self.root: _Node[T] | None = None
        self.size = 0
        for x in items:
            self.insert(x)

    def insert(self, x: T) -> bool:
        """Add x; return False if an item at distance 0 is already stored."""
        if self.root is None:  # @line L1
            self.root = _Node(x)
            self.size = 1
            return True
        t = self.root  # @line L2
        while True:
            k = self.dist(x, t.item)  # @line L3
            if k == 0:  # @line L4
                return False
            child = t.children.get(k)  # @line L5
            if child is None:
                t.children[k] = _Node(x)
                self.size += 1
                return True
            t = child

    def range(self, q: T, r: int) -> list[T]:
        """Return all stored items x with dist(q, x) <= r."""
        out: list[T] = []
        stack = [self.root]
        while stack:
            t = stack.pop()
            if t is None:  # @line L6
                continue
            k = self.dist(q, t.item)  # @line L7
            if k <= r:
                out.append(t.item)
            for i, child in t.children.items():  # @line L8
                if k - r <= i <= k + r:
                    stack.append(child)
        return out

    def nearest(self, q: T) -> tuple[int, T] | None:
        """Return (distance, item) for a stored item nearest to q, or None if the tree is empty."""
        if self.root is None:
            return None
        best: tuple[int, T] | None = None  # @line L9
        stack: list[tuple[int, _Node[T]]] = [(0, self.root)]
        while stack:  # @line L10
            lb, t = stack.pop()
            if best is not None and lb >= best[0]:
                continue
            k = self.dist(q, t.item)  # @line L11
            if best is None or k < best[0]:
                best = (k, t.item)
            for i, child in t.children.items():  # @line L12
                if abs(i - k) < best[0]:
                    stack.append((abs(i - k), child))
        return best  # @line L13