⚡️ Speed up function nms_supercells by 519%

[codeflash-ai]

📄 519% (5.19x) speedup for nms_supercells in unstructured_inference/models/table_postprocess.py

⏱️ Runtime : 148 milliseconds → 23.9 milliseconds (best of 90 runs)

📝 Explanation and details

The optimized code achieves a 518% speedup by addressing the quadratic computational overhead in nms_supercells through two key optimizations:

Primary Optimization: Early Overlap Detection (O(N²) → O(N²) with fast path)

What changed: The optimized code precomputes row_sets and col_sets as Python sets for all supercells upfront, then uses fast set intersection checks (row_set2 & row_sets[supercell1_num]) to skip pairs that cannot possibly overlap before calling the expensive remove_supercell_overlap function.

Why this is faster: In the original code, every pair of supercells (157,238 pairs in the profiler results) called remove_supercell_overlap, which immediately created sets from lists and computed intersections. The line profiler shows this function consumed 1.39s out of 1.52s total (91.6% of runtime).

In the optimized version, the early check filters out 155,637 non-overlapping pairs with cheap set operations, reducing remove_supercell_overlap calls from 157,238 to just 1,601 (a 98% reduction). The function's total time drops from 509ms to 17.9ms. This explains why test cases with many non-overlapping supercells see dramatic improvements:

• test_many_non_overlapping_supercells: 1.24ms → 221μs (459% faster)
• test_large_supercells_list_with_few_overlaps: 19.1ms → 2.70ms (608% faster)
• test_performance_with_moderate_scale: 115ms → 17.5ms (562% faster)

Trade-off: Test cases with dense overlaps (where most pairs actually overlap) show modest slowdowns (10-17%) because they pay the overhead of set precomputation and the additional intersection check without gaining many skips. For example, test_identical_supercells is 10% slower because both supercells overlap completely, so the early check doesn't help.

Secondary Optimization: Reduced Dictionary Lookups

What changed: Inside remove_supercell_overlap, the optimized code stores supercell1/2["row_numbers"] and supercell1/2["column_numbers"] in local variables (rows1, rows2, cols1, cols2) to avoid repeated dictionary key lookups in the while loop.

Why this matters: Dictionary lookups in Python have overhead, and the original code performed them on every loop iteration when checking lengths and calling min()/max(). By caching these references, the optimized version reduces per-iteration overhead, contributing to the ~28x speedup in remove_supercell_overlap (509ms → 17.9ms).

Performance Characteristics

The optimization is highly effective for workloads with sparse overlaps (the common case in Non-Maximum Suppression for real table detection), where most supercell pairs don't overlap. It works well when:

• Many supercells are spatially separated (no shared rows/columns)
• The grid is large relative to supercell size
• Only a small fraction of pairs have actual overlap

It's less beneficial when:

• All supercells overlap the same region (dense overlaps)
• Working with very small supercell counts (<10) where setup overhead dominates
• Single-cell or tiny supercells where overlap checks are already fast

The optimization doesn't change algorithmic complexity but dramatically reduces the constant factor by avoiding expensive operations on non-overlapping pairs—a classic "early exit" pattern that proves very effective for this NMS use case.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 41 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import copy
import random

# imports
import pytest  # used for our unit tests
from unstructured_inference.models.table_postprocess import nms_supercells

def test_non_overlapping_supercells_are_kept_and_sorted():
    # Two supercells that don't overlap at all should both be kept.
    # They should be returned sorted by score (descending).
    a = {"row_numbers": [0, 1], "column_numbers": [0, 1], "score": 0.5, "id": "a"}
    b = {"row_numbers": [2, 3], "column_numbers": [2, 3], "score": 0.9, "id": "b"}
    # Provide in unsorted order to ensure nms_supercells sorts by score internally.
    input_supercells = [a, b]
    codeflash_output = nms_supercells(copy.deepcopy(input_supercells)); result = codeflash_output # 10.4μs -> 10.1μs (3.33% faster)

def test_overlap_removes_lower_confidence_single_cell_supercell():
    # If a lower-confidence supercell overlaps at a single grid cell with a higher-confidence
    # supercell, the lower-confidence one should be shrunk; if it becomes empty it should be removed.
    high = {"row_numbers": [0, 1], "column_numbers": [0, 1], "score": 0.9, "id": "high"}
    low = {"row_numbers": [1], "column_numbers": [1], "score": 0.8, "id": "low"}
    # high has higher score and should remain; low overlaps and should be suppressed
    codeflash_output = nms_supercells([high, low]); res = codeflash_output # 12.3μs -> 14.3μs (14.1% slower)

def test_shrinks_columns_if_rows_fewer_than_columns_and_preserves_if_still_valid():
    # When supercell2 has fewer rows than columns, it will remove overlapping columns first.
    # Construct a scenario where after removal there remain >=2 columns so the supercell is not suppressed.
    high = {"row_numbers": [0], "column_numbers": [0, 1], "score": 0.9, "id": "high"}
    low = {"row_numbers": [0], "column_numbers": [0, 1, 2, 3], "score": 0.8, "id": "low"}
    # Work on a deep copy to avoid modifying references used in assertions
    low_copy = copy.deepcopy(low)
    codeflash_output = nms_supercells([high, low_copy]); res = codeflash_output # 14.1μs -> 16.3μs (13.7% slower)
    # Both should be kept because after removing overlapping columns low_copy will still have 2 columns (2 & 3).
    ids = [r["id"] for r in res]
    # Check that the low supercell was actually shrunk (its columns should have been reduced)
    # Find the low dict in the result
    low_result = next((r for r in res if r["id"] == "low"), None)

def test_single_1x1_supercell_remains_if_first_in_sorted_list():
    # Important nuance: the algorithm only checks suppression for supercells with index >= 1
    # Therefore a single supercell that is 1x1 remains (no suppression check is performed on the highest-scored element).
    single = {"row_numbers": [0], "column_numbers": [0], "score": 0.99, "id": "only_one"}
    codeflash_output = nms_supercells([single]); res = codeflash_output # 5.63μs -> 6.62μs (14.9% slower)

def test_lower_index_1x1_is_suppressed_but_1xN_survives():
    # If a lower-scored supercell is 1x1 it will be suppressed.
    # But a 1xN (one row, multiple columns) should survive unless both dims <2.
    high = {"row_numbers": [0, 1], "column_numbers": [0, 1], "score": 0.95, "id": "high"}
    low_1x1 = {"row_numbers": [2], "column_numbers": [2], "score": 0.5, "id": "low_1x1"}
    low_1xN = {"row_numbers": [3], "column_numbers": [3, 4, 5], "score": 0.4, "id": "low_1xN"}
    codeflash_output = nms_supercells([high, low_1x1, low_1xN]); res = codeflash_output # 13.7μs -> 12.0μs (14.0% faster)
    # high always kept; low_1x1 is suppressed because it is lower indexed (supercell2 case) and 1x1
    ids = [r["id"] for r in res]

def test_middle_row_overlap_causes_entire_supercell_to_be_removed_when_min_max_not_in_common():
    # If the overlap is only at a middle row (neither min nor max), the algorithm sets row_numbers=[]
    # which should cause the supercell to be suppressed.
    high = {"row_numbers": [2], "column_numbers": [2], "score": 0.9, "id": "high"}
    # low has rows [1,2,3] and columns [2,3] so the only common row is 2 which is not min nor max
    low = {"row_numbers": [1, 2, 3], "column_numbers": [2, 3], "score": 0.8, "id": "low"}
    codeflash_output = nms_supercells([high, low]); res = codeflash_output # 12.3μs -> 14.6μs (15.6% slower)

def test_empty_input_returns_empty_list():
    # Passing an empty list should return an empty list quickly and deterministically
    codeflash_output = nms_supercells([]) # 4.99μs -> 5.18μs (3.55% slower)

def test_idempotent_on_larger_random_but_deterministic_dataset():
    # Generate a deterministic pseudo-random set of supercells that create many overlaps.
    # Then verify that applying nms_supercells twice yields the same result (idempotency).
    rnd = random.Random(0)  # deterministic seed
    supercells = []
    num = 40  # moderate "large scale" while keeping runtime reasonable
    max_grid = 12
    for i in range(num):
        # Each supercell gets 1..4 rows and 1..4 columns chosen deterministically
        rcount = rnd.randint(1, 4)
        ccount = rnd.randint(1, 4)
        # pick contiguous ranges for clarity and deterministic overlap patterns
        start_r = rnd.randint(0, max_grid - rcount)
        start_c = rnd.randint(0, max_grid - ccount)
        rows = list(range(start_r, start_r + rcount))
        cols = list(range(start_c, start_c + ccount))
        score = rnd.random()
        supercells.append({"row_numbers": rows, "column_numbers": cols, "score": score, "id": f"s{i}"})

    # Work on deep copies to prevent unexpected in-place mutation between calls
    original = copy.deepcopy(supercells)
    codeflash_output = nms_supercells(copy.deepcopy(original)); first = codeflash_output # 888μs -> 287μs (210% faster)
    # Running again should give the same result (idempotency)
    codeflash_output = nms_supercells(copy.deepcopy(first)); second = codeflash_output # 218μs -> 56.1μs (290% faster)

    # Additional sanity checks: result must be sorted by score descending
    scores = [s["score"] for s in first]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from unstructured_inference.models.table_postprocess import nms_supercells

def test_empty_supercells_list():
    """Test nms_supercells with an empty list of supercells."""
    codeflash_output = nms_supercells([]); result = codeflash_output # 5.09μs -> 5.13μs (0.760% slower)

def test_single_supercell():
    """Test nms_supercells with a single supercell."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1, 2],
            "column_numbers": [0, 1, 2],
        }
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 5.54μs -> 6.70μs (17.2% slower)

def test_two_non_overlapping_supercells():
    """Test nms_supercells with two non-overlapping supercells."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [3, 4],
            "column_numbers": [3, 4],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.4μs -> 10.2μs (2.16% faster)

def test_two_overlapping_supercells_with_shrinking():
    """Test nms_supercells with two overlapping supercells that can be resolved by shrinking."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1, 2],
            "column_numbers": [0, 1, 2],
        },
        {
            "score": 0.7,
            "row_numbers": [1, 2, 3],
            "column_numbers": [1, 2, 3],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 14.5μs -> 17.1μs (15.0% slower)

def test_scores_are_sorted_correctly():
    """Test that nms_supercells sorts by score in descending order."""
    supercells = [
        {
            "score": 0.5,
            "row_numbers": [0],
            "column_numbers": [0],
        },
        {
            "score": 0.95,
            "row_numbers": [1],
            "column_numbers": [1],
        },
        {
            "score": 0.7,
            "row_numbers": [2],
            "column_numbers": [2],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 13.4μs -> 11.8μs (13.7% faster)
    # Check that highest scores are preserved
    scores = [obj["score"] for obj in result]

def test_supercell_with_single_cell():
    """Test nms_supercells with supercells that contain only a single grid cell."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0],
            "column_numbers": [0],
        },
        {
            "score": 0.8,
            "row_numbers": [0],
            "column_numbers": [0],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 12.1μs -> 14.1μs (13.8% slower)

def test_multiple_non_overlapping_supercells():
    """Test nms_supercells with multiple non-overlapping supercells."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [2, 3],
            "column_numbers": [2, 3],
        },
        {
            "score": 0.7,
            "row_numbers": [4, 5],
            "column_numbers": [4, 5],
        },
        {
            "score": 0.6,
            "row_numbers": [6, 7],
            "column_numbers": [6, 7],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 17.6μs -> 13.6μs (29.4% faster)

def test_supercell_with_empty_rows():
    """Test handling of supercells with empty row_numbers list."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [],
            "column_numbers": [0, 1],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.5μs -> 10.0μs (4.48% faster)

def test_supercell_with_empty_columns():
    """Test handling of supercells with empty column_numbers list."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [0, 1],
            "column_numbers": [],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.7μs -> 10.3μs (3.08% faster)

def test_supercell_with_one_row_one_column():
    """Test handling of supercells with exactly one row and one column (1x1 cell)."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [5],
            "column_numbers": [5],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.2μs -> 9.87μs (3.79% faster)

def test_supercell_with_one_row_two_columns():
    """Test handling of supercells with one row and two columns (1x2 cell)."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [5],
            "column_numbers": [5, 6],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.4μs -> 10.1μs (3.93% faster)

def test_supercell_with_two_rows_one_column():
    """Test handling of supercells with two rows and one column (2x1 cell)."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [5, 6],
            "column_numbers": [5],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.5μs -> 10.0μs (4.95% faster)

def test_identical_supercells():
    """Test nms_supercells with two identical supercells (same rows and columns)."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1, 2],
            "column_numbers": [0, 1, 2],
        },
        {
            "score": 0.8,
            "row_numbers": [0, 1, 2],
            "column_numbers": [0, 1, 2],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 16.4μs -> 18.2μs (10.1% slower)

def test_complete_overlap_suppressed():
    """Test that a completely overlapping supercell is suppressed."""
    supercells = [
        {
            "score": 0.95,
            "row_numbers": [0, 1, 2, 3],
            "column_numbers": [0, 1, 2, 3],
        },
        {
            "score": 0.85,
            "row_numbers": [1, 2],
            "column_numbers": [1, 2],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 14.3μs -> 16.6μs (14.0% slower)

def test_partial_row_overlap():
    """Test supercells with partial row overlap."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1, 2],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.7,
            "row_numbers": [2, 3],
            "column_numbers": [0, 1],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 12.7μs -> 15.3μs (17.0% slower)

def test_partial_column_overlap():
    """Test supercells with partial column overlap."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1, 2],
        },
        {
            "score": 0.7,
            "row_numbers": [0, 1],
            "column_numbers": [2, 3],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 13.8μs -> 16.3μs (15.7% slower)

def test_very_high_score_vs_low_score():
    """Test NMS with very high and very low confidence scores."""
    supercells = [
        {
            "score": 0.9999,
            "row_numbers": [0, 1, 2],
            "column_numbers": [0, 1, 2],
        },
        {
            "score": 0.0001,
            "row_numbers": [1, 2, 3],
            "column_numbers": [1, 2, 3],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 14.6μs -> 17.2μs (15.5% slower)

def test_zero_score():
    """Test supercells with zero scores."""
    supercells = [
        {
            "score": 0.8,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.0,
            "row_numbers": [2, 3],
            "column_numbers": [2, 3],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.7μs -> 10.2μs (5.14% faster)

def test_negative_score():
    """Test supercells with negative scores."""
    supercells = [
        {
            "score": 0.5,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": -0.5,
            "row_numbers": [2, 3],
            "column_numbers": [2, 3],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 10.5μs -> 10.1μs (3.92% faster)

def test_equal_scores():
    """Test nms_supercells with supercells having equal scores."""
    supercells = [
        {
            "score": 0.8,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.8,
            "row_numbers": [1, 2],
            "column_numbers": [1, 2],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 12.5μs -> 14.9μs (15.8% slower)

def test_supercell_shrunk_to_empty_rows():
    """Test when a supercell is shrunk completely from rows."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1, 2, 3],
        },
        {
            "score": 0.8,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1, 2, 3],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 16.3μs -> 18.7μs (12.7% slower)

def test_large_row_small_column_supercell():
    """Test supercell with many rows and few columns."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0, 1, 2, 3, 4, 5],
            "column_numbers": [0],
        },
        {
            "score": 0.7,
            "row_numbers": [0, 1, 2, 3, 4, 5],
            "column_numbers": [0, 1],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 18.5μs -> 20.9μs (11.6% slower)

def test_small_row_large_column_supercell():
    """Test supercell with few rows and many columns."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0],
            "column_numbers": [0, 1, 2, 3, 4, 5],
        },
        {
            "score": 0.7,
            "row_numbers": [0, 1],
            "column_numbers": [0, 1, 2, 3, 4, 5],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 17.4μs -> 20.1μs (13.5% slower)

def test_many_non_overlapping_supercells():
    """Test nms_supercells with many non-overlapping supercells."""
    # Create 50 non-overlapping supercells in a grid
    supercells = []
    for i in range(50):
        row_start = (i // 10) * 2
        col_start = (i % 10) * 2
        supercells.append(
            {
                "score": 0.9 - (i * 0.001),
                "row_numbers": [row_start, row_start + 1],
                "column_numbers": [col_start, col_start + 1],
            }
        )
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 1.24ms -> 221μs (459% faster)

def test_chain_overlapping_supercells():
    """Test nms_supercells with a chain of overlapping supercells."""
    # Create a chain where each supercell overlaps with the next
    supercells = []
    for i in range(30):
        supercells.append(
            {
                "score": 0.95 - (i * 0.001),
                "row_numbers": [i, i + 1, i + 2],
                "column_numbers": [0, 1, 2],
            }
        )
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 517μs -> 250μs (107% faster)

def test_grid_overlapping_supercells():
    """Test nms_supercells with a grid of slightly overlapping supercells."""
    # Create a 10x10 grid of supercells with slight overlaps
    supercells = []
    for i in range(10):
        for j in range(10):
            supercells.append(
                {
                    "score": 0.9 - ((i * 10 + j) * 0.0001),
                    "row_numbers": [i * 2, i * 2 + 1, i * 2 + 2],
                    "column_numbers": [j * 2, j * 2 + 1, j * 2 + 2],
                }
            )
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 5.06ms -> 1.24ms (309% faster)

def test_dense_overlapping_supercells():
    """Test nms_supercells with densely packed overlapping supercells."""
    # Create 100 supercells all in the same region with varying scores
    supercells = []
    for i in range(100):
        supercells.append(
            {
                "score": 1.0 - (i * 0.01),
                "row_numbers": [0, 1, 2, 3],
                "column_numbers": [0, 1, 2, 3],
            }
        )
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 4.80ms -> 1.24ms (287% faster)

def test_random_rectangular_supercells():
    """Test nms_supercells with supercells of various rectangular dimensions."""
    supercells = [
        {
            "score": 0.95,
            "row_numbers": [0, 1, 2, 3],
            "column_numbers": [0, 1],
        },
        {
            "score": 0.90,
            "row_numbers": [1, 2],
            "column_numbers": [0, 1, 2, 3, 4],
        },
        {
            "score": 0.85,
            "row_numbers": [5, 6, 7],
            "column_numbers": [5, 6, 7, 8],
        },
        {
            "score": 0.80,
            "row_numbers": [10],
            "column_numbers": [10, 11, 12, 13, 14],
        },
        {
            "score": 0.75,
            "row_numbers": [15, 16, 17, 18, 19],
            "column_numbers": [15],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 29.6μs -> 23.6μs (25.3% faster)

def test_large_supercells_list_with_few_overlaps():
    """Test nms_supercells with 200 supercells where few overlap."""
    supercells = []
    overlap_count = 0
    for i in range(200):
        # Most supercells are non-overlapping, but some are intentionally overlapping
        if i < 180:
            # Non-overlapping grid
            row_idx = (i // 20)
            col_idx = (i % 20)
            supercells.append(
                {
                    "score": 0.95 - (i * 0.0001),
                    "row_numbers": [row_idx * 2, row_idx * 2 + 1],
                    "column_numbers": [col_idx * 2, col_idx * 2 + 1],
                }
            )
        else:
            # Create some overlapping supercells
            overlap_count += 1
            supercells.append(
                {
                    "score": 0.5 - (overlap_count * 0.01),
                    "row_numbers": [0, 1, 2],
                    "column_numbers": [0, 1, 2],
                }
            )
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 19.1ms -> 2.70ms (608% faster)

def test_performance_with_moderate_scale():
    """Test that nms_supercells performs efficiently with moderate scale."""
    # Create 500 supercells with controlled overlaps
    supercells = []
    for i in range(500):
        # Arrange in a way that creates some overlaps but not excessive
        block = i // 50
        pos_in_block = i % 50
        row_start = (block * 3) + (pos_in_block // 10)
        col_start = (pos_in_block % 10) * 2
        supercells.append(
            {
                "score": 0.99 - (i * 0.00001),
                "row_numbers": [row_start, row_start + 1, row_start + 2],
                "column_numbers": [col_start, col_start + 1],
            }
        )
    # This should complete in reasonable time
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 115ms -> 17.5ms (562% faster)

def test_all_supercells_below_threshold():
    """Test when all supercells are below the 2x2 minimum threshold."""
    supercells = [
        {
            "score": 0.9,
            "row_numbers": [0],
            "column_numbers": [0],
        },
        {
            "score": 0.8,
            "row_numbers": [1],
            "column_numbers": [1],
        },
        {
            "score": 0.7,
            "row_numbers": [2],
            "column_numbers": [2],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 15.3μs -> 11.9μs (28.9% faster)

def test_complex_overlapping_pattern():
    """Test nms_supercells with a complex pattern of overlaps."""
    supercells = [
        {
            "score": 0.99,
            "row_numbers": [0, 1, 2, 3],
            "column_numbers": [0, 1, 2, 3],
        },
        {
            "score": 0.95,
            "row_numbers": [2, 3, 4, 5],
            "column_numbers": [2, 3, 4, 5],
        },
        {
            "score": 0.90,
            "row_numbers": [4, 5, 6, 7],
            "column_numbers": [0, 1, 2, 3],
        },
        {
            "score": 0.85,
            "row_numbers": [1, 2],
            "column_numbers": [4, 5, 6, 7],
        },
        {
            "score": 0.80,
            "row_numbers": [6, 7, 8, 9],
            "column_numbers": [4, 5, 6, 7],
        },
    ]
    codeflash_output = nms_supercells(supercells); result = codeflash_output # 32.6μs -> 29.0μs (12.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-nms_supercells-mkouh7ze and push.