~/problems / Heaps / Heaps and priority queues

Balance pins across columns

easy ~15 min Pinterest

A masonry-style feed shows pins in k side-by-side columns, numbered 0 .. k-1 from left to right. Each column grows downward, and its height is the sum of the heights of the pins placed in it (an empty column has height 0).

Pins arrive in the order of heights (pin i is heights[i] pixels tall), and each one is dropped into the column that is currently shortest. If several columns tie for shortest, the pin goes into the leftmost of them.

Write layout(heights: list[int], k: int) -> list[list[int]] that returns, for each column from left to right, the list of pin indices placed in it, top to bottom.

layout([30, 10, 20, 15, 5], 2)   # [[0, 3], [1, 2, 4]]
layout([4, 4, 4], 5)             # [[0], [1], [2], [], []]
layout([], 3)                    # [[], [], []]

Tracing the first call (column heights after each pin):

pin height goes to heights
0 30 column 0 (both 0, leftmost wins) 30, 0
1 10 column 1 30, 10
2 20 column 1 30, 30
3 15 column 0 (both 30, leftmost wins) 45, 30
4 5 column 1 45, 35

Constraints: 1 <= k <= 10**5, len(heights) <= 2 * 10**5, 1 <= heights[i] <= 10**4. Scanning all k columns for every pin is O(n·k), too slow when both are large: aim for O((n + k) log k).

Show hint

Each pin needs the column with the smallest (height, column number). Keep the columns in a structure that hands you that minimum quickly and lets you put the column back after its height grows.

Topic: Heaps and priority queues. heapq: top-k, k-way merge, two heaps for a running median.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc