iterative branch and bound

dedupe/branch_and_bound.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import functools
-import warnings
 from typing import Any, Iterable, Mapping, Sequence, Tuple
 
 from ._typing import Cover
@@ -27,13 +26,11 @@
 def _uncovered_by(
     coverage: Mapping[Any, frozenset[int]], covered: frozenset[int]
 ) -> dict[Any, frozenset[int]]:
-    remaining = {}
-    for predicate, uncovered in coverage.items():
-        still_uncovered = uncovered - covered
-        if still_uncovered:
-            remaining[predicate] = still_uncovered
-
-    return remaining
+    return {
+        pred: still_uncovered
+        for pred, uncovered in coverage.items()
+        if (still_uncovered := uncovered - covered)
+    }
 
 
 def _order_by(
@@ -46,41 +43,22 @@
     return sum(p.cover_count for p in partial)
 
 
-def _suppress_recursion_error(func):
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except RecursionError:
-            warnings.warn("Recursion limit eached while searching for predicates")
-
-    return wrapper
-
-
-def search(candidates, target: int, max_calls: int) -> Partial:
-    calls = max_calls
-
-    cheapest_score = float("inf")
-    cheapest: Partial = ()
-
-    original_cover = candidates.copy()
-
+def search(original_cover: Cover, target: int, calls: int) -> Partial:
     def _covered(partial: Partial) -> int:
         return (
             len(frozenset.union(*(original_cover[p] for p in partial)))
             if partial
             else 0
         )
 
-    @_suppress_recursion_error
-    def walk(candidates: Cover, partial: Partial = ()) -> None:
-        nonlocal calls
-        nonlocal cheapest
-        nonlocal cheapest_score
+    cheapest_score = float("inf")
+    cheapest: Partial = ()
 
-        if calls <= 0:
-            return
+    start: tuple[Cover, Partial] = (original_cover, ())
+    to_explore = [start]
 
-        calls -= 1
+    while to_explore and calls:
+        candidates, partial = to_explore.pop()
 
         covered = _covered(partial)
         score = _score(partial)
@@ -97,17 +75,16 @@
                 order_by = functools.partial(_order_by, candidates)
                 best = max(candidates, key=order_by)
 
-                remaining = _uncovered_by(candidates, candidates[best])
-                walk(remaining, partial + (best,))
-                del remaining
-
                 reduced = _remove_dominated(candidates, best)
-                walk(reduced, partial)
-                del reduced
+                to_explore.append((reduced, partial))
+
+                remaining = _uncovered_by(candidates, candidates[best])
+                to_explore.append((remaining, partial + (best,)))
 
         elif score < cheapest_score:
             cheapest = partial
             cheapest_score = score
 
-    walk(candidates)
+        calls -= 1
+
     return cheapest