Fix suboptimal assignment in cell-matching Hungarian algorithm

marimo/_utils/cell_matching.py

@@ -79,82 +79,100 @@ def pop_local(available: list[tuple[int, CellId_t]], idx: int) -> CellId_t:
     return available.pop(best_idx)[1]
 
 
-def _hungarian_algorithm(scores: list[list[float]]) -> list[int]:
-    """Implements the Hungarian algorithm to find the best matching.
-
-    In general this class of problem is known as the assignment problem and is
-    pretty well studied. This is a textbook implementation to avoid additional
-    dependencies. Links:
-    - https://en.wikipedia.org/wiki/Hungarian_algorithm
-    """
-    score_matrix = [row[:] for row in scores]
-    n = len(score_matrix)
-
-    # Step 1: Subtract row minima
-    for i in range(n):
-        min_value = min(score_matrix[i])
-        for j in range(n):
-            score_matrix[i][j] -= min_value
-
-    # Step 2: Subtract column minima
-    for j in range(n):
-        min_value = min(score_matrix[i][j] for i in range(n))
+# Above this size the exact O(n^3) solver gets slow on dense, tie-heavy cost
+# matrices -- in particular the zero-padded matrices produced when many more
+# cells are added than removed (~0.5s at n=500 for a realistic matrix, several
+# seconds for the padded worst case). Such large simultaneous edits are rare and
+# a slightly sub-optimal match there is harmless, so fall back to a fast O(n^2)
+# greedy assignment above the cutoff.
+_MAX_OPTIMAL_ASSIGNMENT_SIZE = 100
+
+
+def _greedy_assignment(scores: list[list[float]]) -> list[int]:
+    """Fast approximate assignment; `result[column] = row`, same convention as
+    `_hungarian_algorithm`."""
+    n = len(scores)
+    result = [-1] * n
+    used_row = [False] * n
+    # Assign the most decisive columns (smallest best cost) first.
+    for j in sorted(
+        range(n), key=lambda c: min(scores[r][c] for r in range(n))
+    ):
+        best_row, best_cost = -1, float("inf")
         for i in range(n):
-            score_matrix[i][j] -= min_value
-
-    # Step 3: Find initial assignment
-    row_assignment = [-1] * n
-    col_assignment = [-1] * n
-
-    # Find independent zeros
-    for i in range(n):
-        for j in range(n):
-            if (
-                score_matrix[i][j] == 0
-                and row_assignment[i] == -1
-                and col_assignment[j] == -1
-            ):
-                row_assignment[i] = j
-                col_assignment[j] = i
-
-    # Step 4: Improve assignment iteratively
-    while True:
-        assigned_count = sum(1 for x in row_assignment if x != -1)
-        if assigned_count == n:
-            break
+            if not used_row[i] and scores[i][j] < best_cost:
+                best_cost, best_row = scores[i][j], i
+        if best_row != -1:
+            used_row[best_row] = True
+            result[j] = best_row
+    return result
 
-        # Find minimum uncovered value
-        min_uncovered = float("inf")
-        for i in range(n):
-            for j in range(n):
-                if row_assignment[i] == -1 and col_assignment[j] == -1:
-                    min_uncovered = min(min_uncovered, score_matrix[i][j])
 
-        if min_uncovered == float("inf"):
-            break
+def _hungarian_algorithm(scores: list[list[float]]) -> list[int]:
+    """Solve the assignment problem, returning a minimum-cost matching.
 
-        # Update matrix
-        for i in range(n):
-            for j in range(n):
-                if row_assignment[i] == -1 and col_assignment[j] == -1:
-                    score_matrix[i][j] -= min_uncovered
-                elif row_assignment[i] != -1 and col_assignment[j] != -1:
-                    score_matrix[i][j] += min_uncovered
+    Uses the O(n^3) shortest-augmenting-path method (Jonker-Volgenant /
+    Kuhn-Munkres), which is guaranteed to find an optimal assignment without
+    additional dependencies. Links:
+    - https://en.wikipedia.org/wiki/Hungarian_algorithm
 
-        # Try to find new assignments
-        for i in range(n):
-            if row_assignment[i] == -1:
-                for j in range(n):
-                    if score_matrix[i][j] == 0 and col_assignment[j] == -1:
-                        row_assignment[i] = j
-                        col_assignment[j] = i
-                        break
-
-    # Convert to result format
+    Returns a list `result` where `result[column] = row` for the row matched to
+    each column (or -1 if unmatched, which only happens for an empty input).
+    """
+    n = len(scores)
+    if n == 0:
+        return []
+
+    inf = float("inf")
+    # Potentials (u for rows, v for columns) and the current column -> row
+    # matching. Index 0 is a sentinel used while growing the augmenting path,
+    # so everything is 1-indexed.
+    u = [0.0] * (n + 1)
+    v = [0.0] * (n + 1)
+    match_col_to_row = [0] * (n + 1)
+    way = [0] * (n + 1)
+
+    for i in range(1, n + 1):
+        match_col_to_row[0] = i
+        j0 = 0
+        min_val = [inf] * (n + 1)
+        used = [False] * (n + 1)
+        # Grow an alternating tree until we reach an unmatched column.
+        while True:
+            used[j0] = True
+            i0 = match_col_to_row[j0]
+            delta = inf
+            j1 = -1
+            for j in range(1, n + 1):
+                if not used[j]:
+                    cur = scores[i0 - 1][j - 1] - u[i0] - v[j]
+                    if cur < min_val[j]:
+                        min_val[j] = cur
+                        way[j] = j0
+                    if min_val[j] < delta:
+                        delta = min_val[j]
+                        j1 = j
+            # Update potentials so the reduced costs stay non-negative.
+            for j in range(n + 1):
+                if used[j]:
+                    u[match_col_to_row[j]] += delta
+                    v[j] -= delta
+                else:
+                    min_val[j] -= delta
+            j0 = j1
+            if match_col_to_row[j0] == 0:
+                break
+        # Augment along the path recorded in `way`.
+        while j0:
+            j1 = way[j0]
+            match_col_to_row[j0] = match_col_to_row[j1]
+            j0 = j1
+
+    # Convert to result format: result[column] = row (0-indexed).
     result = [-1] * n
-    for i in range(n):
-        if row_assignment[i] != -1:
-            result[row_assignment[i]] = i
+    for j in range(1, n + 1):
+        if match_col_to_row[j] != 0:
+            result[j - 1] = match_col_to_row[j] - 1
 
     return result
 
@@ -251,8 +269,13 @@ def filter_and_backfill() -> list[CellId_t]:
                     # NB. transposed indices for Hungarian
                     scores[y][x] = score
 
-    # Use Hungarian algorithm to find the best matching
-    matches = _hungarian_algorithm(scores)
+    # Use the exact assignment for small problems, and a fast greedy fallback
+    # for large ones where the exact O(n^3) solver would be too slow.
+    matches = (
+        _greedy_assignment(scores)
+        if n > _MAX_OPTIMAL_ASSIGNMENT_SIZE
+        else _hungarian_algorithm(scores)
+    )
     for idx, code in enumerate(next_codes):
         if result[idx] is None:
             match_idx = next_order[next_inverse[code]].pop(0)

tests/_utils/test_cell_matching.py

@@ -0,0 +1,127 @@
+# Copyright 2026 Marimo. All rights reserved.
+from __future__ import annotations
+
+import itertools
+import random
+
+from marimo._utils.cell_matching import (
+    _greedy_assignment,
+    _hungarian_algorithm,
+    match_cell_ids_by_similarity,
+)
+
+
+def _assignment_cost(scores: list[list[float]], result: list[int]) -> float:
+    """Total cost of the matching returned by _hungarian_algorithm.
+
+    `result[column] = row`; raises if the matching is not a permutation.
+    """
+    n = len(scores)
+    col_to_row = {j: result[j] for j in range(n) if result[j] != -1}
+    assert len(col_to_row) == n, "matching is not complete"
+    assert len(set(col_to_row.values())) == n, "matching is not a permutation"
+    return sum(scores[row][col] for col, row in col_to_row.items())
+
+
+def _brute_force_optimal(scores: list[list[float]]) -> float:
+    n = len(scores)
+    return min(
+        sum(scores[i][perm[i]] for i in range(n))
+        for perm in itertools.permutations(range(n))
+    )
+
+
+def test_hungarian_empty() -> None:
+    assert _hungarian_algorithm([]) == []
+
+
+def test_hungarian_single() -> None:
+    assert _hungarian_algorithm([[5.0]]) == [0]
+
+
+def test_hungarian_known_suboptimal_case() -> None:
+    # Regression test: the previous covering heuristic returned a cost-17
+    # assignment here; the optimal cost is 13.
+    scores = [
+        [7.0, 7.0, 8.0],
+        [3.0, 5.0, 3.0],
+        [3.0, 7.0, 4.0],
+    ]
+    result = _hungarian_algorithm(scores)
+    assert _assignment_cost(scores, result) == 13.0
+    assert _assignment_cost(scores, result) == _brute_force_optimal(scores)
+
+
+def test_hungarian_matches_brute_force() -> None:
+    # The assignment must be optimal for every matrix, not merely valid.
+    rng = random.Random(20260825)
+    for _ in range(500):
+        n = rng.randint(1, 6)
+        scores = [
+            [float(rng.randint(0, 9)) for _ in range(n)] for _ in range(n)
+        ]
+        result = _hungarian_algorithm([row[:] for row in scores])
+        assert _assignment_cost(scores, result) == _brute_force_optimal(scores)
+
+
+def test_hungarian_handles_negative_and_float_costs() -> None:
+    rng = random.Random(1234)
+    for _ in range(200):
+        n = rng.randint(1, 5)
+        scores = [[rng.uniform(-5.0, 5.0) for _ in range(n)] for _ in range(n)]
+        result = _hungarian_algorithm([row[:] for row in scores])
+        assert (
+            abs(
+                _assignment_cost(scores, result) - _brute_force_optimal(scores)
+            )
+            < 1e-9
+        )
+
+
+def test_greedy_assignment_returns_valid_permutation() -> None:
+    rng = random.Random(7)
+    for _ in range(50):
+        n = rng.randint(1, 25)
+        scores = [[rng.uniform(-5.0, 5.0) for _ in range(n)] for _ in range(n)]
+        result = _greedy_assignment([row[:] for row in scores])
+        assert sorted(result) == list(range(n))
+
+
+def test_greedy_assignment_scales_to_large_inputs() -> None:
+    # The exact O(n^3) solver is too slow on large, tie-heavy padded matrices
+    # (several seconds at n=500); the greedy fallback used above the size cutoff
+    # must stay fast and still return a valid assignment.
+    n = 500
+    scores = [[0.0] * n for _ in range(n)]
+    for i in range(n):
+        for j in range(20):
+            scores[i][j] = float((i * 31 + j) % 97)
+    result = _greedy_assignment(scores)
+    assert sorted(result) == list(range(n))
+
+
+def test_match_cell_ids_identical_notebook() -> None:
+    data = {"a": "x = 1", "b": "y = 2", "c": "z = 3"}
+    assert match_cell_ids_by_similarity(dict(data), dict(data)) == {
+        "a": "a",
+        "b": "b",
+        "c": "c",
+    }
+
+
+def test_match_cell_ids_prefers_most_similar() -> None:
+    # Every cell was edited (no exact matches), so matching falls back to the
+    # similarity assignment. Each next cell should keep the id of the prev cell
+    # it most closely resembles.
+    prev = {
+        "imp": "import pandas as pd",
+        "tot": "x = compute_total(data)",
+        "plt": "df.plot(kind='bar')",
+    }
+    nxt = {
+        "n_plt": "df.plot(kind='line')",
+        "n_tot": "x = compute_total(rows)",
+        "n_imp": "import polars as pd",
+    }
+    mapping = match_cell_ids_by_similarity(dict(prev), dict(nxt))
+    assert mapping == {"plt": "n_plt", "tot": "n_tot", "imp": "n_imp"}