Fix suboptimal assignment in cell-matching Hungarian algorithm - #10654
Fix suboptimal assignment in cell-matching Hungarian algorithm#10654winklemad wants to merge 3 commits into
Conversation
_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.
|
I have read the CLA Document and I hereby sign the CLA |
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
All contributors have signed the CLA ✍️ ✅ |
There was a problem hiding this comment.
Pull request overview
Fixes suboptimal cell matching by replacing the flawed assignment algorithm with an optimal Hungarian implementation.
Changes:
- Implements O(n³) minimum-cost assignment matching.
- Adds regression, randomized optimality, edge-case, and integration tests.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
tests/_utils/test_cell_matching.py |
Adds comprehensive matching correctness and integration tests. |
marimo/_utils/cell_matching.py |
Replaces the assignment algorithm while preserving its API. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
dmadisetti
left a comment
dmadisetti
left a comment
There was a problem hiding this comment.
Thanks, but can you bench this on a large notebook? We may be taking a large perf hit. Even if the algorithm is not perfect, I rather a "good but fast minima" vs a perfect minima that's slow
| 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 | ||
| ) |
There was a problem hiding this comment.
If this takes more than a second, I'd like to reduce this
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.
|
Good call — benchmarked it, and you're right about the worst case. The Hungarian only runs on the changed cell subset (
So on ordinary edits it's within a few ms of the old code, but a large bulk reload could freeze for seconds — exactly the "perfect but slow" case you want to avoid. Fixed in a8c590a: added a fast O(n²) greedy assignment and a size cutoff ( On your other comment (the tests at 66–77): measured them — If you'd rather keep it dead simple and just always use greedy (dropping the exact solver entirely), happy to do that instead — let me know. |
for more information, see https://pre-commit.ci
This pull request was authored by a coding agent.
📝 Summary
_hungarian_algorithminmarimo/_utils/cell_matching.pydoes not solve the assignment problem correctly. Its covering step (steps 3–4) treats "row/column currently assigned" as a proxy for the Hungarian line cover, which is not the correct covering rule, so it frequently settles on a non-minimum assignment.Measured against a brute-force optimum on random 2×2–6×6 integer cost matrices, it returns a suboptimal assignment about a third of the time. Minimal example:
Why it matters
match_cell_ids_by_similaritybuilds asimilarity_scorecost matrix between deleted and added cells and uses this to pair them (lower score = more similar). It's called fromcell_manager(andcompiler/lint) to preserve cell ids when a notebook is edited/reloaded. When the matching is suboptimal, an edited cell is paired with a less-similar previous cell than it should be, so a cell's identity — and with it its outputs, UI element state, and reactive bindings — can be transferred to the wrong cell.Fix
Replace the body of
_hungarian_algorithmwith the O(n³) shortest-augmenting-path method (Jonker–Volgenant / Kuhn–Munkres), which is guaranteed to return a minimum-cost matching and needs no new dependencies. The signature and theresult[column] = rowreturn convention are unchanged, so all callers are unaffected. Handles empty/1×1 inputs, ties, negative and floating-point costs.Tests
Added
tests/_utils/test_cell_matching.py:match_cell_ids_by_similaritymapping an unchanged notebook to itself and pairing edited cells with their most similar predecessors.All three optimality tests fail on the previous implementation and pass with the fix.
tests/_ast/test_cell_manager.pycontinues to pass (60 passed, 1 pre-existing xfail).ruff format/ruff check/mypyclean on the changed files.📋 Pre-Review Checklist
✅ Merge Checklist