CF 102431E - Non-Maximum Suppression

Rating: -
Tags: -
Solve time: 6m 59s
Verified: yes

Solution

Problem Understanding

Each detection is a square of the same side length S. Its position is determined by the bottom-left corner (x, y), and it has a distinct confidence score. NMS processes these detections from highest score to lowest score. A detection is selected if it has not already been suppressed. Once selected, it suppresses every lower-scoring detection whose intersection-over-union is strictly greater than the given threshold. The task is to output the indices of all detections that survive this process, sorted by their original indices.

The official constraints allow n to reach 10^5, with S up to 10^7, and the threshold is an exact value with three decimal places. The current Codeforces statement gives a 15 second time limit and 256 MB memory limit. An O(n^2) implementation can perform about 10^10 pair checks when n = 10^5, which is far beyond what is reasonable even with the relatively generous time limit. We need the number of comparisons per selected box to be bounded by a constant.

For two equal squares, let

dx = |x1 - x2| and dy = |y1 - y2|.

Their intersection has width max(0, S - dx) and height max(0, S - dy). If the intersection area is A, the union area is 2S^2 - A, so the IoU is

A / (2S^2 - A).

We should not evaluate this using floating point. If the threshold is p / 1000, then

A / (2S^2 - A) > p / 1000

is exactly equivalent to

(1000 + p) A > 2000 p S^2.

Every quantity in this inequality is an integer.

Several edge cases can silently break a careless implementation. First, the comparison is strictly greater than the threshold. For example,

The intersection area is 3 * 2 = 6, the union is 18 - 6 = 12, and the IoU is exactly 0.500. The correct output is

Case #1: 2
1 2

Using >= threshold would incorrectly suppress box 2.

A second edge case is zero overlap. For example,

The squares only touch at their boundaries, so their intersection area is zero. Both detections must be selected:

Case #1: 2
1 2

A careless implementation that reasons from distance alone can accidentally treat touching squares as overlapping.

A third edge case is identical boxes. For example,

The highest-scoring box is selected first, and its IoU with both other boxes is 1, which is strictly greater than 0.700. The answer is

Case #1: 1
3

The distinct-score guarantee tells us exactly which identical box survives.

Finally, the selected detections must be printed in increasing index order, not in score order. If boxes 3 and 1 survive while box 2 is suppressed, the output must be 1 3, even though box 3 might have the larger score. The official sample demonstrates this ordering.

Approaches

The direct solution follows NMS literally. Sort all boxes by decreasing score, then maintain the currently unsuppressed boxes. When the next box is selected, compare it against every remaining box and suppress those with IoU above the threshold. This is correct because it is exactly the definition of the NMS process.

The problem is the number of comparisons. In the worst case, no box suppresses another, so the first selected box checks roughly n boxes, the second checks roughly n - 1, and so on. The total is about n(n - 1)/2, which reaches roughly 5 * 10^9 comparisons for n = 10^5. That is already too large before accounting for sorting, hash-table operations, and geometry calculations.

The key observation is that all squares have exactly the same size. For a fixed selected square, another square can only have IoU above the threshold when its bottom-left corner is close enough to the selected corner. In the (x, y) plane, the set of possible conflicting positions is a bounded region around the selected point.

We can turn this geometric locality into a grid. Choose a cell size C so that any two points in the same cell are guaranteed to have IoU strictly greater than the threshold. Then at most one box from each cell can ever be selected. We can also show that two conflicting boxes must lie in cells whose coordinates differ by at most two. Consequently, when processing a selected box, we only need to inspect the 5 x 5 cells surrounding its cell.

The cell size is chosen exactly rather than approximately. Let the threshold be p / 1000 and define

q = 2p / (1000 + p).

Two boxes have IoU greater than the threshold exactly when their intersection area is greater than qS^2.

Suppose two bottom-left corners are in the same cell of side C. Their coordinate differences are at most C - 1, so their intersection area is at least

(S - C + 1)^2.

We choose the largest integer C satisfying

(1000 + p)(S - C + 1)^2 > 2000pS^2.

Thus every pair of boxes in the same cell conflicts. This is the central density bound.

Now consider two boxes that really conflict. Since

(S - dx)(S - dy) > qS^2,

both factors must individually be larger than sqrt(q) S. Hence

dx < S(1 - sqrt(q))

and similarly for dy.

The chosen cell size is at least S(1 - sqrt(q)), so a conflicting pair differs by less than twice one cell width in either coordinate. Their grid coordinates can consequently differ by at most two. The 5 x 5 neighborhood contains every possible suppressible box.

The grid does not need to physically delete suppressed boxes. A box belongs to one cell, and a cell can contain at most one selected box because all boxes in that cell conflict with each other. Therefore any particular stored box can be examined only by selected boxes from at most 25 neighboring cells. Each box participates in only a constant number of comparisons overall.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n²)	O(n)	Too slow
Optimal	O(n log n)	O(n)	Accepted

The official contest material also describes the same geometric idea: equal-sized squares only conflict inside a fixed neighborhood, allowing a grid-based search instead of checking every box.

Algorithm Walkthrough

Parse the threshold as an integer p representing p / 1000. Parse every score as an integer scaled by 10^6. This avoids floating-point comparisons when sorting and when deciding whether IoU is strictly above the threshold.
Sort the box indices by decreasing score. Since all scores are distinct, this gives exactly the order in which NMS considers boxes.
Compute the largest grid-cell side length C satisfying

(1000 + p)(S - C + 1)^2 > 2000pS^2.

The left-hand side decreases as C increases, so C can be found with binary search. The strict inequality is deliberate. It matches the strict IoU > threshold rule. 4. Insert every box into a dictionary keyed by (x // C, y // C). The dictionary stores the indices belonging to each spatial cell. 5. Process the boxes in descending score order. If a box has already been suppressed, skip it. Otherwise select it and mark it as processed. 6. For a selected box, inspect all grid cells with offsets from -2 through 2 in both coordinates. Every box that could have IoU above the threshold must be in one of these 25 cells. 7. For every candidate that has not already been suppressed, compute its exact intersection area and test

(1000 + p) * intersection > 2000 * p * S^2.

If the inequality holds, mark that candidate as suppressed. 8. Store every selected index, then sort those indices increasingly before printing. The NMS process itself is score-ordered, while the required output is index-ordered.

Why it works

The invariant is that before processing a box in score order, every earlier box is either selected or suppressed exactly as prescribed by NMS. If the current box is already suppressed, it must never be selected. Otherwise, no previously selected box has suppressed it, so it is exactly the highest-scoring remaining box and must be selected.

When a box is selected, every lower-scoring box that it can possibly suppress lies in the inspected 5 x 5 neighborhood. The cell-size construction guarantees that every same-cell pair conflicts, while the overlap inequality bounds every conflicting pair to cells at distance at most two. The exact integer IoU test then suppresses precisely those candidates whose IoU is strictly larger than the threshold. Thus every NMS decision is made correctly, and the final selected set is exactly the required one.

Python Solution

import sys
input = sys.stdin.readline

def parse_scaled(s, digits):
    if '.' in s:
        a, b = s.split('.')
    else:
        a, b = s, ''
    b = (b + '0' * digits)[:digits]
    return int(a) * (10 ** digits) + int(b)

def solve():
    t = int(input())
    output = []

    for case_id in range(1, t + 1):
        n, S, threshold = input().split()
        n = int(n)
        S = int(S)

        # threshold is exact to three decimal places.
        p = parse_scaled(threshold, 3)

        boxes = []
        order = []

        for i in range(n):
            x, y, score = input().split()
            x = int(x)
            y = int(y)
            score = parse_scaled(score, 6)
            boxes.append((x, y))
            order.append((score, i))

        order.sort(reverse=True)

        # We need the largest integer C such that
        #
        # (1000 + p) * (S - C + 1)^2 > 2000 * p * S^2
        #
        # Then any two boxes in one cell necessarily have IoU > threshold.
        target = 2000 * p * S * S
        coefficient = 1000 + p

        def good(c):
            overlap = S - c + 1
            return coefficient * overlap * overlap > target

        lo, hi = 1, S
        while lo < hi:
            mid = (lo + hi + 1) // 2
            if good(mid):
                lo = mid
            else:
                hi = mid - 1

        cell_size = lo

        grid = {}
        for _, idx in order:
            x, y = boxes[idx]
            key = (x // cell_size, y // cell_size)
            grid.setdefault(key, []).append(idx)

        suppressed = bytearray(n)
        selected = []

        for _, idx in order:
            if suppressed[idx]:
                continue

            suppressed[idx] = 1
            selected.append(idx + 1)

            x1, y1 = boxes[idx]
            gx = x1 // cell_size
            gy = y1 // cell_size

            for ox in range(-2, 3):
                for oy in range(-2, 3):
                    candidates = grid.get((gx + ox, gy + oy))
                    if candidates is None:
                        continue

                    for j in candidates:
                        if suppressed[j]:
                            continue

                        x2, y2 = boxes[j]

                        ix = S - abs(x1 - x2)
                        if ix <= 0:
                            continue

                        iy = S - abs(y1 - y2)
                        if iy <= 0:
                            continue

                        area = ix * iy

                        if coefficient * area > target:
                            suppressed[j] = 1

        selected.sort()

        output.append(f"Case #{case_id}: {len(selected)}")
        output.append(" ".join(map(str, selected)))

    sys.stdout.write("\n".join(output))

if __name__ == "__main__":
    solve()

The input parsing deliberately keeps the threshold and scores as integers. A threshold such as 0.500 becomes 500, while a score such as 0.900000 becomes 900000. Since the input guarantees exact decimal values, this representation preserves their ordering and avoids rounding errors.

The binary search finds the largest safe cell size. For a candidate cell size C, the smallest possible intersection between two boxes in that cell occurs when their two coordinates differ by C - 1. The resulting intersection is (S - C + 1)^2, so the good predicate directly checks whether even that worst case still has IoU strictly above the threshold.

The grid dictionary maps each spatial cell to all original box indices in that cell. We keep suppressed entries instead of deleting them. This keeps the implementation simple, and it does not destroy the complexity bound. Since at most one box from any cell can survive, a stored box can be encountered only from the constant number of selected boxes whose cells are within distance two.

The suppressed bytearray records both already selected boxes and boxes eliminated by NMS. Python's bytearray is substantially smaller than a list of Python booleans for 10^5 entries.

The intersection calculation uses

ix = S - abs(x1 - x2)

and the analogous expression for iy. A non-positive value means there is no positive-area intersection, so the IoU is zero and the candidate cannot be suppressed. The final comparison uses integer arithmetic and preserves the strict inequality exactly.

Python integers have arbitrary precision, so products such as S^2 and 2000 * p * S^2 do not overflow. In languages with fixed-width integers, a 64-bit integer is sufficient for these bounds.

Worked Examples

The provided sample has three squares of side 4 and threshold 0.390. The boxes are already listed in descending score order.

Step	Box	Cell	Suppressed before	IoU checks	Action
1	1	`(0, 0)`	none	box 2, box 3	Select 1, suppress 2
2	2	`(0, 0)`	yes	none	Skip
3	3	`(1, 1)`	no	box 1	Select 3

For the first and second boxes, the coordinate differences are (1, 1), giving intersection area 9. The union is 32 - 9 = 23, so their IoU is 9/23, which is greater than 0.390. For boxes 1 and 3, the intersection area is 4, the union is 28, and the IoU is 1/7, which is below the threshold. The selected indices are consequently 1 and 3.

A second example exercises the exact-threshold boundary.

1
3 3 0.500
0 0 0.900
0 1 0.800
10 10 0.700

For S = 3 and threshold 0.500, the first two boxes have intersection area 6 and union area 12.

Step	Box	Cell	IoU with selected boxes	Action
1	1	`(0, 0)`	none	Select 1
2	2	`(0, 1)`	`6 / 12 = 0.500`	Select 2
3	3	`(10, 10)`	`0`	Select 3

The second box survives because the rule is strictly greater than the threshold. The final result is

Case #1: 3
1 2 3

This trace confirms that the integer comparison does not accidentally turn equality into suppression.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log n)	Sorting costs O(n log n), while every box participates in only O(1) spatial comparisons
Space	O(n)	Boxes, sorting indices, suppression state, and grid storage all use O(n) memory

The grid construction requires a binary search over S, adding only O(log S) work per test case. Since S <= 10^7, this is fewer than 24 iterations. The dominant cost is sorting 10^5 boxes, followed by a constant number of spatial checks per box. The official limits are n <= 10^5, S <= 10^7, 15 seconds, and 256 MB, so the resulting O(n log n) solution fits the intended scale.

Test Cases

The following test harness contains the solution as a callable function so that every assertion can be executed directly. The maximum-size case uses 100000 identical positions with distinct scores, forcing NMS to retain exactly the highest-scoring box.

import io
import sys

def parse_scaled(s, digits):
    if '.' in s:
        a, b = s.split('.')
    else:
        a, b = s, ''
    b = (b + '0' * digits)[:digits]
    return int(a) * (10 ** digits) + int(b)

def solution(data: str) -> str:
    it = iter(data.split())
    t = int(next(it))
    output = []

    for case_id in range(1, t + 1):
        n = int(next(it))
        S = int(next(it))
        threshold = next(it)
        p = parse_scaled(threshold, 3)

        boxes = []
        order = []

        for i in range(n):
            x = int(next(it))
            y = int(next(it))
            score = parse_scaled(next(it), 6)
            boxes.append((x, y))
            order.append((score, i))

        order.sort(reverse=True)

        target = 2000 * p * S * S
        coefficient = 1000 + p

        def good(c):
            overlap = S - c + 1
            return coefficient * overlap * overlap > target

        lo, hi = 1, S
        while lo < hi:
            mid = (lo + hi + 1) // 2
            if good(mid):
                lo = mid
            else:
                hi = mid - 1

        cell_size = lo

        grid = {}
        for _, idx in order:
            x, y = boxes[idx]
            key = (x // cell_size, y // cell_size)
            grid.setdefault(key, []).append(idx)

        suppressed = bytearray(n)
        selected = []

        for _, idx in order:
            if suppressed[idx]:
                continue

            suppressed[idx] = 1
            selected.append(idx + 1)

            x1, y1 = boxes[idx]
            gx = x1 // cell_size
            gy = y1 // cell_size

            for ox in range(-2, 3):
                for oy in range(-2, 3):
                    candidates = grid.get((gx + ox, gy + oy))
                    if candidates is None:
                        continue

                    for j in candidates:
                        if suppressed[j]:
                            continue

                        x2, y2 = boxes[j]

                        ix = S - abs(x1 - x2)
                        if ix <= 0:
                            continue

                        iy = S - abs(y1 - y2)
                        if iy <= 0:
                            continue

                        area = ix * iy

                        if coefficient * area > target:
                            suppressed[j] = 1

        selected.sort()

        output.append(f"Case #{case_id}: {len(selected)}")
        output.append(" ".join(map(str, selected)))

    return "\n".join(output)

def run(inp: str) -> str:
    return solution(inp)

sample = """\
1
3 4 0.390
0 0 0.9
1 1 0.8
2 2 0.7
"""

assert run(sample) == """\
Case #1: 2
1 3
""", "provided sample"

boundary = """\
1
3 3 0.500
0 0 0.900
0 1 0.800
10 10 0.700
"""

assert run(boundary) == """\
Case #1: 3
1 2 3
""", "exact threshold must not suppress"

minimum = """\
1
1 1 0.700
0 0 1.000
"""

assert run(minimum) == """\
Case #1: 1
1
""", "minimum-size input"

identical = """\
1
5 5 0.300
2 2 0.100000
2 2 0.900000
2 2 0.500000
2 2 0.700000
2 2 0.300000
"""

assert run(identical) == """\
Case #1: 1
2
""", "identical boxes keep only the highest score"

far_apart = """\
1
4 10 0.700
0 0 0.400000
10 0 0.900000
0 10 0.800000
10 10 0.700000
"""

assert run(far_apart) == """\
Case #1: 4
1 2 3 4
""", "zero-overlap boxes all survive"

n = 100000
lines = [f"{n} 1 0.700"]
for i in range(n):
    score = (n - i) / 100000
    lines.append(f"0 0 {score:.5f}")

maximum = "\n".join(["1"] + lines) + "\n"
maximum_output = run(maximum)
maximum_lines = maximum_output.splitlines()

assert maximum_lines[0] == "Case #1: 1", "maximum-size count"
assert maximum_lines[1] == "1", "maximum-size highest score survives"

print("All tests passed.")

Test input	Expected output	What it validates
Provided sample	`Case #1: 2`, indices `1 3`	Basic NMS behavior and output ordering
`S=3`, threshold `0.500`, IoU exactly `0.500`	All three selected	Strict `>` boundary
`n=1`, `S=1`	Only index `1`	Minimum-size input
Five identical boxes	Only the highest-scoring index	Identical positions and score ordering
Four completely separated boxes	All four selected	Zero-overlap handling
`n=100000`, `S=1`	Only index `1`	Maximum `n`, dense identical positions, and scalability

Edge Cases

For the strict-threshold case,

the first box is selected. The second box has intersection 6 and union 12, so the IoU is exactly 0.5. The integer test becomes an equality:

(1000 + 500) * 6 = 2000 * 500 * 3^2 / 12,

or more directly,

1500 * 6 = 500 * 12.

The condition uses > rather than >=, so box 2 is not suppressed. The algorithm selects both indices.

For zero overlap,

the horizontal difference is 4, equal to the side length. Thus ix = 4 - 4 = 0, and the implementation immediately skips the candidate. No floating-point IoU calculation is necessary. Both boxes survive.

For identical boxes,

all three points belong to the same grid cell. The first processed box is index 3 because it has the largest score. Its intersection with every other box is 1, so the exact suppression condition succeeds because IoU is 1 > 0.7. Both remaining boxes are marked suppressed and never selected.

For output ordering, consider

The boxes are sufficiently separated that all survive. NMS processes them in the order 2, 3, 1, but the required output is

Case #1: 3
1 2 3

The final sort of selected handles this distinction between processing order and output order.

The maximum-size case uses 100000 boxes at the same coordinate. They all occupy one grid cell, and the cell construction guarantees that boxes in that cell have IoU above the threshold. The highest-scoring box is selected first and suppresses the other 99999 boxes. Every later entry is skipped. The work remains linear apart from the initial sort because all those boxes are contained in one spatial bucket and each is examined only a constant number of times by selected boxes.