Design a frequency-ordered container with weighted random pick

01The design, in the six-step framework

1. Requirements and use cases

A 2026 new-grad round: build an object that organises items by how often they were added, then a follow-up that asks for a random item where popular items come up more often. The second part changes the data structure, and the interviewer wants to hear you notice.

• A caller adds an item; its count rises by one.
• A caller asks for the count of an item in O(1).
• A caller asks for the most frequent item in O(1), and the top k items in O(k log n) or better.
• A caller removes one occurrence of an item.
• Follow-up: a weighted random pick where an item with count c is chosen with probability c / total.

Out of scope: decay over time, persistence, concurrency unless asked. Non-functional constraints: add, count and remove in O(1); most-frequent in O(1); top-k without sorting everything; and the weighted pick in O(log n) with a structure that tolerates counts changing.

2. Core entities and relationships

• FrequencyContainer: the public object; owns counts, frequency buckets and the maximum frequency; answers the ordering questions.
• Bucket: the set of items that currently share one count, kept in a dictionary from count to an ordered set so ties resolve by recency.
• WeightedPicker: the follow-up collaborator; a Fenwick tree over item slots, where each slot's weight is the item's count.

The container composes its buckets and, when the follow-up arrives, composes a picker and keeps it in step on every add and remove. Buckets indexed by count are the same device as in the LFU cache: a count only moves by one, so an item moves between neighbouring buckets, and the maximum moves predictably. The picker is a separate class because weighted sampling has nothing to do with ordering; keeping it apart lets the first half of the interview end with a complete, small object (single responsibility).

3. Interfaces and key methods

import random
from typing import Generic, Hashable, List, Optional, TypeVar

T = TypeVar("T", bound=Hashable)


class FrequencyContainer(Generic[T]):
    def __init__(self, rng: Optional[random.Random] = None) -> None: ...
    def add(self, item: T) -> None:
        """count += 1 in O(1); keeps buckets, max_count and the picker in step."""
    def remove(self, item: T) -> bool:
        """count -= 1; False if the item is absent; drops the item at zero."""
    def count(self, item: T) -> int: ...
    def most_frequent(self) -> Optional[T]:
        """O(1): the most recently added item in the highest bucket."""
    def top(self, k: int) -> List[T]:
        """Up to k items by count descending, ties by recency."""
    def random_weighted(self) -> T:
        """Item with probability count / total; raises IndexError when empty."""

Invariants: counts[x] == c exactly when x is in bucket c; max_count is the largest non-empty bucket, or zero; the picker's weight for x equals counts[x]; and no item has count zero in any structure.

4. Design patterns and the choices behind them

Buckets by count for the ordering half

The force is O(1) for add and for most-frequent. A heap gives O(log n) add and cannot update an arbitrary item's priority without an index. Buckets keyed by count give O(1) moves because the count changes by exactly one each time; the maximum goes up by one on an add, and on a remove falls to c − 1 only when the top bucket empties. Heaps are the right tool for top-k once, not for maintaining a live ranking under unit changes.

A Fenwick tree for the weighted pick

The force is sampling proportional to counts that keep changing. Prefix sums over counts plus binary search give O(log n) picks but O(n) updates when a count changes. A Fenwick tree keeps prefix sums under point updates, both in O(log n): draw a uniform integer u in [0, total), descend the tree to the first slot whose prefix sum exceeds u. The naive alternative of expanding items into a list with one entry per count is O(1) to pick but O(count) to remove and unbounded in memory; it is worth mentioning as the first idea and rejecting out loud.

Recency as the tie-break, kept in the bucket

The force is "which of two items with equal counts is first". An ordered dictionary per bucket gives the most recently added item at the end in O(1), the same trick as the LRU structure. State the rule; some interviewers prefer the oldest, and the only change is which end you read.

5. Core flow and code

Flow for random_weighted() with counts a:3, b:1 (total 4): the picker draws u in {0,1,2,3}; prefix sums are a:3, b:4; u in {0,1,2} lands on a, u = 3 on b, giving the 3:1 ratio required.

1. Draw u = rng.randrange(total).
2. Walk the Fenwick tree from the top power of two downward, moving right whenever the subtree sum is at most u, subtracting as you go.
3. The slot reached is the item; map it back through item_at.

import random
from collections import OrderedDict
from typing import Dict, Generic, Hashable, List, Optional, TypeVar

T = TypeVar("T", bound=Hashable)


class WeightedPicker(Generic[T]):
    """Fenwick tree over item slots; weight = count. Slots are never freed (weights go to 0)."""

    def __init__(self, capacity_hint: int = 16) -> None:
        self.tree: List[int] = [0] * (capacity_hint + 1)
        self.slot_of: Dict[T, int] = {}
        self.item_at: List[Optional[T]] = [None] * (capacity_hint + 1)
        self.n = 0

    def _grow(self) -> None:
        old_n = len(self.tree) - 1
        weights = [self.prefix(i) - self.prefix(i - 1) for i in range(1, old_n + 1)]
        self.tree = [0] * (2 * old_n + 1)
        self.item_at += [None] * old_n
        for i, w in enumerate(weights, start=1):
            self._update(i, w)

    def _update(self, i: int, delta: int) -> None:
        while i < len(self.tree):
            self.tree[i] += delta
            i += i & -i

    def prefix(self, i: int) -> int:
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s

    def adjust(self, item: T, delta: int) -> None:
        slot = self.slot_of.get(item)
        if slot is None:
            if self.n + 1 >= len(self.tree):
                self._grow()
            self.n += 1
            slot = self.n
            self.slot_of[item] = slot
            self.item_at[slot] = item
        self._update(slot, delta)

    def total(self) -> int:
        return self.prefix(self.n)

    def pick(self, u: int) -> T:
        """Smallest slot whose prefix sum exceeds u (0 <= u < total)."""
        pos, step = 0, 1
        while step * 2 <= self.n:
            step *= 2
        while step:
            nxt = pos + step
            if nxt <= self.n and self.tree[nxt] <= u:
                pos = nxt
                u -= self.tree[nxt]
            step //= 2
        return self.item_at[pos + 1]  # type: ignore[return-value]


class FrequencyContainer(Generic[T]):
    def __init__(self, rng: Optional[random.Random] = None) -> None:
        self.counts: Dict[T, int] = {}
        self.buckets: Dict[int, "OrderedDict[T, None]"] = {}
        self.max_count = 0
        self.picker: WeightedPicker[T] = WeightedPicker()
        self.rng = rng or random.Random()

    def add(self, item: T) -> None:
        c = self.counts.get(item, 0)
        if c:
            del self.buckets[c][item]
            if not self.buckets[c]:
                del self.buckets[c]
        self.counts[item] = c + 1
        self.buckets.setdefault(c + 1, OrderedDict())[item] = None
        self.max_count = max(self.max_count, c + 1)
        self.picker.adjust(item, +1)

    def remove(self, item: T) -> bool:
        c = self.counts.get(item)
        if not c:
            return False
        del self.buckets[c][item]
        if not self.buckets[c]:
            del self.buckets[c]
            if self.max_count == c:
                self.max_count = c - 1
        if c == 1:
            del self.counts[item]
        else:
            self.counts[item] = c - 1
            self.buckets.setdefault(c - 1, OrderedDict())[item] = None
        self.picker.adjust(item, -1)
        return True

    def count(self, item: T) -> int:
        return self.counts.get(item, 0)

    def most_frequent(self) -> Optional[T]:
        if not self.max_count:
            return None
        return next(reversed(self.buckets[self.max_count]))

    def top(self, k: int) -> List[T]:
        out: List[T] = []
        for c in sorted(self.buckets, reverse=True):
            for item in reversed(self.buckets[c]):
                if len(out) == k:
                    return out
                out.append(item)
        return out

    def random_weighted(self) -> T:
        total = self.picker.total()
        if total == 0:
            raise IndexError("empty container")
        return self.picker.pick(self.rng.randrange(total))


if __name__ == "__main__":
    fc: FrequencyContainer[str] = FrequencyContainer(random.Random(3))
    for x in ["a", "a", "a", "b"]:
        fc.add(x)
    assert fc.most_frequent() == "a" and fc.top(2) == ["a", "b"] and fc.count("a") == 3
    draws = [fc.random_weighted() for _ in range(4000)]
    assert 0.70 < draws.count("a") / 4000 < 0.80          # expected 0.75
    assert fc.remove("a") and fc.remove("a") and fc.remove("a") and fc.count("a") == 0
    assert fc.most_frequent() == "b" and fc.random_weighted() == "b"

6. Extensions, concurrency and what interviewers probe

• Notice the structure change. The first half needs buckets; the weighted pick needs prefix sums. Saying "the follow-up needs a different structure, here is why" before coding is what the round rewards.
• Top-k complexity. Walking buckets from the top is O(k + number of distinct counts visited); for a hot top-k with a huge item set, a bounded heap over bucket heads is the next step.
• Slots are never freed. An item that drops to zero keeps its slot with weight zero; it is never picked and costs one integer. If churn is high, compact the picker when zero-weight slots exceed half.
• Randomness is injected so the ratio test above is reproducible; see designing for unit tests.
• Thread safety. One lock; the buckets, counts and picker must change together.
• Common mistake. Using sets per bucket (no recency order, so most-frequent among ties is arbitrary) or forgetting to lower max_count when the top bucket empties on remove.

02Where this question was reported

• Prepfully · Prepfully: among recently reported Google coding questions
• LeetCode Discuss · Related LeetCode 384 Shuffle an Array, Google-tagged
• FastPrep · HashMap-based object organising items by frequency; follow-up: return an item randomly weighted by frequency