Place books on the least-full allowed shelf

Every book asks one question about a prefix of the shelves: which shelf in
[0, limit - 1] holds the fewest books, breaking ties by the smallest index?
Then it adds one book to that shelf. So the problem is a stream of
"prefix argmin, then point increment" operations.

Brute force. For each book, scan shelves 0 .. limit - 1, keep the
first shelf with the smallest count, and increment it. This is
O(n * m) for n books and m shelves. It is the right thing to say
first and it makes a good reference for testing, but with n, m up to
10^5 it is too slow.

Optimal: a segment tree over (count, index) pairs. Store at each leaf
the pair (count[i], i). Each internal node stores the minimum of its
children under tuple order. Tuple order compares counts first and indices
second, so the minimum pair is exactly "fewest books, then smallest index".
That turns the tie-break rule into ordinary min and nothing else has to
handle it.

• Query: the minimum pair over [0, limit - 1] is O(log m).
• Update: set the leaf to (count + 1, i) and recompute the parents, also
  O(log m).

An iterative (bottom-up) segment tree keeps the code short. Pad the size up
to a power of two and fill the padding leaves with (inf, inf) so they never
win a comparison.

Invariants.

• tree[size + i] == (count[i], i) for every real shelf i.
• tree[v] == min(tree[2v], tree[2v + 1]) for every internal node v.
• Padding leaves are (inf, inf).

Validation. A limit outside 1 .. numShelves has no valid shelf. Raise
an error (or clamp, if the interviewer prefers). Say which you choose.

Why not a heap? A single heap gives the global minimum, but the range
changes with every book. You would have to pop shelves outside the prefix
and push them back, which is O(m) per book in the worst case. The segment
tree answers any prefix directly.

from __future__ import annotations

import math
from typing import List, Tuple

INF = math.inf


class ShelfTree:
    """Min segment tree over (book_count, shelf_index) pairs."""

    def __init__(self, num_shelves: int) -> None:
        if num_shelves <= 0:
            raise ValueError("need at least one shelf")
        self.n = num_shelves
        size = 1
        while size < num_shelves:
            size *= 2
        self.size = size
        self.tree: List[Tuple[float, float]] = [(INF, INF)] * (2 * size)
        for i in range(num_shelves):
            self.tree[size + i] = (0, i)
        for v in range(size - 1, 0, -1):
            self.tree[v] = min(self.tree[2 * v], self.tree[2 * v + 1])

    def query(self, lo: int, hi: int) -> Tuple[float, float]:
        """Minimum (count, index) over shelves lo..hi inclusive."""
        best: Tuple[float, float] = (INF, INF)
        lo += self.size
        hi += self.size + 1  # half-open [lo, hi)
        while lo < hi:
            if lo & 1:
                best = min(best, self.tree[lo])
                lo += 1
            if hi & 1:
                hi -= 1
                best = min(best, self.tree[hi])
            lo //= 2
            hi //= 2
        return best

    def add(self, shelf: int, delta: int) -> None:
        v = self.size + shelf
        count, idx = self.tree[v]
        self.tree[v] = (count + delta, idx)
        v //= 2
        while v:
            self.tree[v] = min(self.tree[2 * v], self.tree[2 * v + 1])
            v //= 2


def organize_books(num_shelves: int, limits: List[int]) -> List[int]:
    """Return the shelf chosen for each book, in arrival order."""
    tree = ShelfTree(num_shelves)
    placed: List[int] = []
    for limit in limits:
        if not 1 <= limit <= num_shelves:
            raise ValueError(f"limit {limit} outside 1..{num_shelves}")
        _, shelf = tree.query(0, limit - 1)
        shelf = int(shelf)
        tree.add(shelf, 1)
        placed.append(shelf)
    return placed


def organize_books_brute(num_shelves: int, limits: List[int]) -> List[int]:
    counts = [0] * num_shelves
    placed: List[int] = []
    for limit in limits:
        shelf = min(range(limit), key=lambda i: (counts[i], i))
        counts[shelf] += 1
        placed.append(shelf)
    return placed


if __name__ == "__main__":
    import random

    assert organize_books(3, [3, 3, 3, 3]) == [0, 1, 2, 0]
    assert organize_books(4, [1, 2, 2, 4, 4]) == [0, 1, 0, 2, 3]
    assert organize_books(1, [1, 1, 1]) == [0, 0, 0]
    assert organize_books(5, []) == []
    assert organize_books(5, [2, 2, 2, 2]) == [0, 1, 0, 1]
    for _ in range(300):
        m = random.randint(1, 12)
        ls = [random.randint(1, m) for _ in range(random.randint(0, 40))]
        assert organize_books(m, ls) == organize_books_brute(m, ls)
    print("ok")

Let m be the number of shelves and n the number of books.

• Build: O(m) time, O(m) space (the tree has at most 4m nodes).
• Each book: one prefix query and one point update, O(log m) each.
• Total: O(m + n log m) time, O(m + n) space including the output.
• Brute force: O(n * m) time, O(m) space.

• num_shelves = 3, limits = [3, 3, 3, 3] gives [0, 1, 2, 0]. With the
  full range, books go round-robin.
• num_shelves = 4, limits = [1, 2, 2, 4, 4] gives [0, 1, 0, 2, 3]. The
  third book sees counts (1, 1) in range 0..1 and takes shelf 0 by index.
• num_shelves = 5, limits = [2, 2, 2, 2] gives [0, 1, 0, 1]. Shelves
  2..4 are never reachable and stay empty.
• num_shelves = 1, limits = [1, 1, 1] gives [0, 0, 0].
• num_shelves = 5, limits = [] gives [].
• num_shelves = 3, limits = [3, 1, 1, 3] gives [0, 0, 0, 1]. Shelf 0
  piles up under limit = 1; the last book takes shelf 1, not shelf 2.
• limits contains 0 or a value above num_shelves: raise ValueError.
• Randomized: compare against the brute force on small inputs.

1. Return the final counts, or the count on the fullest shelf. Keep a
   plain counts array next to the tree (or read the leaves). The fullest
   shelf is a max query, so add a second tree or track the running max,
   since counts only grow.
2. Arbitrary ranges [lo, hi] instead of prefixes. The segment tree
   already answers any range; only the query bounds change. A prefix-only
   structure (such as a Fenwick tree for prefix minimums) would not carry
   over, which is a good reason to choose the segment tree up front.
3. Books are removed (checkouts). A removal is add(shelf, -1). Point
   updates in both directions cost O(log m). A trick that relies on counts
   only growing (such as keeping a running max) breaks, and the candidate
   should say so.
4. Each shelf has a capacity. When a shelf reaches capacity, set its
   leaf to (inf, index) so it never wins. If the whole range is full, the
   query returns inf; report the book as unplaced or raise. To prefer the
   shelf with the most free space instead, store (-free, index).
5. Books have widths and the rule is "least total width". Store total
   width instead of count and add the width in the update. Nothing else
   changes. If widths are large, keep counts as integers to avoid float
   drift.

Mid-level (L4)

• Restates the tie-break rule and writes the brute force correctly.
• Sees that n * m is too slow at 10^5 and reaches for a segment tree,
  perhaps with hints.
• Has a working tree with correct bounds and a few hand-checked tests.

Senior (L5)

• Uses (count, index) pairs so the tie-break falls out of tuple min,
  and explains why a heap does not fit a moving range.
• Writes a clean iterative or recursive tree without off-by-one errors,
  handles invalid limits, and tests against the brute force.
• Handles removals and capacities quickly by changing leaf values.

Staff (L6+)

• Chooses the structure for the follow-ups up front (ranges, removals,
  capacities) and names what each extension would break.
• Discusses scale: one tree per library branch, concurrency when several
  clerks shelve at once (serialize updates or shard by shelf range), and
  persisting counts so a restart does not lose state.
• Keeps the API small and testable, and states the invariants out loud.