Fix suboptimal assignment in cell-matching Hungarian algorithm

marimo/_utils/cell_matching.py

@@ -80,81 +80,70 @@ def pop_local(available: list[tuple[int, CellId_t]], idx: int) -> CellId_t:
 
 
 def _hungarian_algorithm(scores: list[list[float]]) -> list[int]:
-    """Implements the Hungarian algorithm to find the best matching.
+    """Solve the assignment problem, returning a minimum-cost 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:
+    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
-    """
-    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))
-        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
-
-        # 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
+    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
 
-        # 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
-
-        # 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
+    # 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
 

tests/_utils/test_cell_matching.py

@@ -0,0 +1,104 @@
+# Copyright 2026 Marimo. All rights reserved.
+from __future__ import annotations
+
+import itertools
+import random
+
+from marimo._utils.cell_matching import (
+    _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_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"}