Search the Answer

What is this?

None of these three problems hands you a sorted array with a target in it. What each does have is a quantity where the answer is a boundary: below some side length every square fits, above it none do; below some start position the window is improvable, above it it is not; below some cumulative sum the draw belongs to an earlier index.

Recognising the boundary is the work. The loop that finds it is the same five lines every time.

💡 Fun fact: Find K Closest Elements is usually attempted by finding the closest value and expanding outward with two pointers — O(log n + k) and perfectly fine. The sharper version binary-searches the window's starting index over the range [0, n − k], comparing x − arr[mid] against arr[mid + k] − x: if the element just past the window is closer than the one at its start, the window must shift right. It searches over positions rather than values, which is the idea worth taking away.

🔓 The 3 problems in this chapter are free. Sign in with Google or Microsoft to start solving.

The one-line idea: identify the quantity whose predicate flips once — a size, a position, a running total — and binary-search that. The thing being searched is rarely the thing you were handed.

1. Search a size, verify with a prefix sum

Build the 2-D prefix sum once, padded with a zero row and column:

P = [[0] * (n + 1) for _ in range(m + 1)]
for r in range(m):
    for c in range(n):
        P[r+1][c+1] = mat[r][c] + P[r][c+1] + P[r+1][c] - P[r][c]

def square_sum(r, c, side):                       # top-left (r, c), 0-indexed
    return P[r+side][c+side] - P[r][c+side] - P[r+side][c] + P[r][c]

The - P[r][c] in the build and the + P[r][c] in the query are the same inclusion–exclusion correction: the top-left region is subtracted twice and must be added back once. Padding is what removes every boundary check — without it, r-1 and c-1 need guarding in four places.

With square_sum costing O(1), feasibility for a given side is one sweep of the grid, and the side length is binary-searched over [0, min(m, n)].

2. Search a position

lo, hi = 0, len(arr) - k                          # candidate window starts
while lo < hi:
    mid = (lo + hi) // 2
    if x - arr[mid] > arr[mid + k] - x:           # the far end is closer → shift right
        lo = mid + 1
    else:
        hi = mid
return arr[lo:lo + k]

The comparison is between the element at the window's start and the element just past its end. If dropping the start and gaining the next element is an improvement, the window belongs further right. The > rather than >= is what implements the stated tie-break — on equal distance the smaller value wins, which means keeping the left window.

3. Search a distribution

prefix = list(itertools.accumulate(weights))
target = random.random() * prefix[-1]
return bisect.bisect_right(prefix, target)        # first cumulative sum > target

Each index owns a slice of the line whose length is its weight, so a uniform draw lands in it exactly in proportion. Using bisect_right on a strictly increasing prefix (weights are positive) gives the first index whose cumulative sum exceeds the draw — and drawing from [0, total) rather than (0, total] keeps zero-probability edge cases out.

4. A 30-second worked example (window start)

arr = [1, 2, 3, 4, 5], k = 4, x = 3. Candidate starts are 0 and 1.

lo=0, hi=1        mid = 0
   x - arr[0]     = 3 - 1 = 2
   arr[0+4] - x   = 5 - 3 = 2
   2 > 2 is false → hi = 0
lo == hi == 0     window = arr[0:4] = [1, 2, 3, 4]   ✓

The tie is what makes this instructive: both ends are two away, and the specification prefers the smaller values, so the window stays left. Using >= here would shift it right and return [2,3,4,5] — the wrong answer, and one that passes every non-tied test.

5. Where you'll actually meet this

• Computer vision. Integral images make region queries O(1), and largest-region-under-threshold is this exact search.
• Load balancing and sharding. Weighted server selection per request is the cumulative-sum draw.
• Recommendation and experimentation. Weighted sampling and A/B bucket assignment use the same construction.
• Search UIs. "K nearest results to this value" backs range and date pickers.
• Capacity planning. Smallest size that satisfies a measurable-but-not-invertible constraint.

6. Problems in this chapter

▶ Maximum Side Length of a Square with Sum ≤ Threshold

Largest square whose sum stays under a threshold. Binary-search the side length; check feasibility with a padded 2-D prefix sum.
Pattern: search a size + O(1) region verification. Target: O(mn log(min(m,n))) time, O(mn) space.

▶ Find K Closest Elements

The k values closest to x, in ascending order. Binary-search the window's start over [0, n−k], comparing the element at the start against the one just past the end.
Pattern: search a position. Target: O(log(n−k) + k) time, O(1) extra space.

▶ Random Pick with Weight

Draw an index with probability proportional to its weight. Build cumulative sums once, then binary-search a uniform draw.
Pattern: search a distribution. Target: O(n) build, O(log n) per pick.

7. Common pitfalls 🚫

• Getting the inclusion–exclusion sign wrong. The corner is subtracted twice and must be added back once, in both the build and the query.
• Not padding the prefix table, which forces boundary checks in four places and invites off-by-one errors.
• Using >= in the window comparison. The tie-break says prefer the smaller values, which means keeping the left window.
• Searching over [0, n] for the window start. The range is [0, n − k], or the slice runs off the end.
• Drawing from (0, total] in the weighted pick, which can exclude the first index when its weight is small.
• Rebuilding the cumulative sums per draw. Build once in the constructor.
• bisect_left instead of bisect_right, which mis-assigns draws that land exactly on a boundary.

8. Key takeaways

1. Binary search needs a predicate that flips once — not an array and not a target.
2. Search sizes, positions and totals, not only values.
3. A padded 2-D prefix sum makes any rectangle query O(1) and removes all the boundary cases.
4. Compare the two ends of a candidate window to decide which way it must move.
5. Cumulative sums turn a uniform draw into a weighted one — binary search used to generate rather than to find.
6. Tie-breaks are part of the specification, and here one comparison operator implements the whole rule.
7. Why interviewers like it: none of these announce themselves as binary search, so spotting the monotonic boundary is the entire signal.

Order: Maximum Side Length of a Square with Sum ≤ Threshold → Find K Closest Elements → Random Pick with Weight.