From f3855b92838b2e9545513c15c24e9b3096aae880 Mon Sep 17 00:00:00 2001 From: Madan Kumar Date: Tue, 25 Aug 2026 17:10:24 +0530 Subject: [PATCH 1/3] Fix suboptimal assignment in cell-matching Hungarian algorithm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _hungarian_algorithm used a covering step that treated 'currently assigned' as a proxy for the Hungarian line cover, which is incorrect. On random 2x2-6x6 cost matrices it returned a non-minimum assignment about a third of the time (e.g. [[7,7,8],[3,5,3],[3,7,4]] gave cost 17 instead of the optimal 13). match_cell_ids_by_similarity relies on this to pair edited/added cells with the most similar previous cells (used when preserving cell ids across notebook edits in cell_manager). A suboptimal matching transfers a cell's identity — and its outputs, UI state, and reactive bindings — to the wrong cell. Replace the body with the O(n^3) shortest-augmenting-path method (Jonker-Volgenant / Kuhn-Munkres), which always finds a minimum-cost matching, with no new dependencies. The signature and return convention are unchanged. Adds tests asserting optimality against a brute-force baseline. --- marimo/_utils/cell_matching.py | 129 +++++++++++++---------------- tests/_utils/test_cell_matching.py | 104 +++++++++++++++++++++++ 2 files changed, 163 insertions(+), 70 deletions(-) create mode 100644 tests/_utils/test_cell_matching.py diff --git a/marimo/_utils/cell_matching.py b/marimo/_utils/cell_matching.py index 459f1303cd5..c6b0013b25d 100644 --- a/marimo/_utils/cell_matching.py +++ b/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 diff --git a/tests/_utils/test_cell_matching.py b/tests/_utils/test_cell_matching.py new file mode 100644 index 00000000000..4453375d212 --- /dev/null +++ b/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"} From a8c590ab290b3896e885613249479bdb8b36a591 Mon Sep 17 00:00:00 2001 From: Madan Kumar Date: Fri, 28 Aug 2026 09:29:04 +0530 Subject: [PATCH 2/3] Bound worst-case assignment time with a greedy fallback above n=100 The exact O(n^3) solver is fine on realistic cell-change matrices (~10ms at 100 cells, ~75ms at 200) but can take several seconds on the large, tie-heavy zero-padded matrices produced when many more cells are added than removed (~8s at n=500). Add a fast O(n^2) greedy assignment and use it above a size cutoff, keeping the exact optimum for the common small case while capping the worst case to a few tens of ms. Adds tests for the greedy path. --- marimo/_utils/cell_matching.py | 36 ++++++++++++++++++++++++++++-- tests/_utils/test_cell_matching.py | 23 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/marimo/_utils/cell_matching.py b/marimo/_utils/cell_matching.py index c6b0013b25d..b4ad45a4e45 100644 --- a/marimo/_utils/cell_matching.py +++ b/marimo/_utils/cell_matching.py @@ -79,6 +79,33 @@ def pop_local(available: list[tuple[int, CellId_t]], idx: int) -> CellId_t: return available.pop(best_idx)[1] +# 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): + 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 + + def _hungarian_algorithm(scores: list[list[float]]) -> list[int]: """Solve the assignment problem, returning a minimum-cost matching. @@ -240,8 +267,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) diff --git a/tests/_utils/test_cell_matching.py b/tests/_utils/test_cell_matching.py index 4453375d212..53542fc6863 100644 --- a/tests/_utils/test_cell_matching.py +++ b/tests/_utils/test_cell_matching.py @@ -5,6 +5,7 @@ import random from marimo._utils.cell_matching import ( + _greedy_assignment, _hungarian_algorithm, match_cell_ids_by_similarity, ) @@ -77,6 +78,28 @@ def test_hungarian_handles_negative_and_float_costs() -> None: ) +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)) == { From 179f5045677c2be1676b99067e0919089545080d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:59:35 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- marimo/_utils/cell_matching.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/marimo/_utils/cell_matching.py b/marimo/_utils/cell_matching.py index b4ad45a4e45..dda07ea184e 100644 --- a/marimo/_utils/cell_matching.py +++ b/marimo/_utils/cell_matching.py @@ -95,7 +95,9 @@ def _greedy_assignment(scores: list[list[float]]) -> list[int]: 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))): + 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): if not used_row[i] and scores[i][j] < best_cost: