Weighted Random & Hashmap Design

Affirm's coding loop keeps circling back to one problem, and once you have seen it you cannot unsee it. Build a small container. It has to insert a value and delete a value in constant time, and it has to hand back a random member in constant time too. That is the whole brief. The container shows up in a different costume nearly every round. One interviewer calls it a set that tolerates duplicates. The next calls it a dictionary that has to forget a random key. The costume keeps changing. The machine underneath does not.

The tension is the fun part. A plain list gives you an O(1) random pick but an O(n) delete, because deleting from the middle shifts everything behind it. A plain hash set flips the tradeoff: delete is O(1), but there is no way to pick a uniformly random key without walking the whole thing. Affirm wants both halves at once. Getting there takes two structures that cover for each other, plus one small move called swap-to-tail.

The Affirm angle

Affirm reskins this one problem more than any other in its loop, so learning to recognize the reskins on sight is most of the battle. The prompt always arrives under-specified on purpose. Does the container hold duplicates or not. Are the values integers or strings. Should the random pick be uniform over distinct values or weighted toward the ones inserted most. The interviewer is watching whether you ask before you build, because a correct solution to the wrong reading of the prompt is how strong candidates still get cut here. Pin down the variant first. Then the data structures almost pick themselves.

Every version Affirm asks is the same core wearing a new label. Each one that comes up changes a single thing.

• Insert, delete, get random with no duplicates. The textbook version. One value maps to one slot, and get_random picks a uniform slot.
• Duplicates allowed. A value can appear many times, so the map points each value to a set of slots rather than a single index. get_random now favors frequent values for free, since a value with five slots is five times as likely to be picked.
• String keys and values. Nothing structural changes. The map is keyed by the value you store, and a string hashes just as well as an int.
• deleteRandom in O(1). Evict a random element rather than a named one. Pick a random slot, then run the same swap-to-tail delete on it.
• Weighted by frequency. get_random must return a value with probability proportional to how many times it was inserted. The duplicate design already does this, which is why the duplicates variant and the weighted variant are secretly the same problem.
• getRandom fast, delete slow. A weaker cousin where put and get and getRandom are O(1) but delete is O(n). Recognizing it tells you that swap-to-tail is the exact trick that upgrades that delete to O(1).

Start with the version that carries every twist inside it: duplicates allowed, values possibly repeated. The array is a flat list of every occurrence, one slot per copy. The map sends each value to the set of slots that currently hold it. Insert is the easy half. Append the value to the array, and record its new slot in that value's set. Both moves are O(1).

from collections import defaultdict
class RandomizedCollection:
    def __init__(self):
        self.vals = []                 # one slot per occurrence, in insertion order
        self.idx = defaultdict(set)    # value -> the set of slots that hold it
    def insert(self, val):
        was_absent = len(self.idx[val]) == 0
        self.idx[val].add(len(self.vals))   # its new slot is the tail index
        self.vals.append(val)
        return was_absent                   # True only when the value is newly present

Delete is where the trick lives. You cannot pull a value out of the middle of the array without shifting everything behind it, which would cost O(n). So you do not delete in place. You take any slot that holds the target value, copy the array's last element into that slot, and shrink the array by one. The hole is filled from the tail, so nothing shifts. The only bookkeeping left is to tell the map that the moved element lives in its new slot now, and that the old tail slot is gone.

    def remove(self, val):
        if not self.idx[val]:
            return False                     # no copy present: nothing to remove
        i = self.idx[val].pop()              # take one slot holding val
        last = self.vals[-1]
        last_i = len(self.vals) - 1
        self.vals[i] = last                  # move the tail value into the freed slot
        self.idx[last].add(i)                # the moved value now lives at slot i
        self.idx[last].discard(last_i)       # and no longer at the old tail slot
        self.vals.pop()                      # drop the tail
        return True

The random pick is one line, and it is O(1). Choose a uniform slot in the array and return whatever value sits there. Because every copy of a value owns its own slot, a value inserted five times fills five slots, so a uniform pick returns it five times as often as a value inserted once. That is the weighted behavior Affirm asks for, and it costs nothing extra. You get frequency-weighted sampling as a side effect of storing duplicates.

    def get_random(self):
        # every copy owns a slot, so a uniform slot pick is already
        # weighted by how many times each value was inserted
        return random.choice(self.vals)

One caveat. This weights by insertion count only. Say Affirm hands you arbitrary weights that are not counts, like a value that should be twice as likely yet appears once. The O(1) array pick no longer covers it. You move to a prefix sum of weights with a binary search at O(log n) per pick. For the frequency version they usually ask, the array stays O(1).

Build the container end to end. The four deterministic operations are below, plus the get_random that ties the structure together. get_random returns a random value, so its output cannot be pinned to a fixed answer, which is why the checks exercise the deterministic calls instead: insert, remove, contains, count. Implement get_random anyway. It is the point of the whole exercise, and an interviewer will ask you to walk through why it is both O(1) and correctly weighted.

1

Random Access With Duplicates

MediumSignature problem

Problem

Implement RandomizedCollection, a container that allows duplicate values and keeps every operation O(1) on average.

RandomizedCollection() builds an empty collection.
insert(val) adds a copy of val and returns true if val was not already present, or false if it was. The copy is added either way.
remove(val) removes one copy of val and returns true, or returns false and changes nothing if no copy exists.
contains(val) returns true if at least one copy of val is present.
count(val) returns how many copies of val are present.
get_random() returns a random copy from the collection. Each stored copy is equally likely, so a value present k times comes back with probability proportional to k. This is Affirm's weighted twist: values inserted more often are returned more often.

On the random method. get_random returns a random element, so its result cannot be checked against a fixed expected value. The cases below exercise the deterministic methods only. Implement get_random regardless, in O(1), and be ready to explain why it is correctly weighted.

Input encoding. Each test case is a JSON array of operations. The first is ["RandomizedCollection", []]; the rest are ["insert", [val]], ["remove", [val]], ["contains", [val]], and ["count", [val]]. Every call records its return value, and the constructor records null.

Input: RandomizedCollection(); insert(3); insert(3); insert(9); count(3); remove(3); count(3)

Output: [null, true, false, true, 2, true, 1]

The first `insert(3)` adds a new value and returns true. The second `insert(3)` keeps the duplicate and returns false, since a copy of 3 was already there. With two 3s and one 9 in the collection, `count(3)` is 2, and a `get_random()` there would return 3 twice as often as 9, because 3 owns two of the three slots. `remove(3)` then deletes a single copy and returns true, so `count(3)` drops to 1.

How this uses the pattern

Hold an array of every copy and a hashmap from value to the set of array slots that copy occupies. insert appends and records the slot. remove uses swap-to-tail: overwrite one of the value's slots with the last element, fix that element's slot in the map, then pop. get_random picks a uniform slot, which is automatically weighted by how many copies each value has.

Random Access With Duplicates

Python

Focus Mode

Loading...

OutputSubmissions

Run

Submit

Click "Run" to execute your code

Test Cases

Input: [["RandomizedCollection",[]],["insert",[1]],["insert",[1]],["insert",[2]],["count",[1]],["count",[2]]]Expected: [null,true,false,true,2,1]

Input: [["RandomizedCollection",[]],["remove",[5]],["insert",[5]],["contains",[5]],["remove",[5]],["contains",[5]]]Expected: [null,false,true,true,true,false]

Input: [["RandomizedCollection",[]],["insert",[7]],["insert",[7]],["insert",[7]],["count",[7]],["remove",[7]],["count",[7]]]Expected: [null,true,false,false,3,true,2]

Input: [["RandomizedCollection",[]],["insert",[1]],["insert",[2]],["insert",[3]],["remove",[1]],["count",[1]],["count",[3]],["contains",[3]],["count",[2]]]Expected: [null,true,true,true,true,0,1,true,1]

Input: [["RandomizedCollection",[]],["insert",[9]],["insert",[9]],["remove",[9]],["contains",[9]],["remove",[9]],["contains",[9]],["remove",[9]]]Expected: [null,true,false,true,true,true,false,false]

Input: [["RandomizedCollection",[]],["insert",[10]],["insert",[20]],["insert",[10]],["insert",[30]],["count",[10]],["remove",[20]],["contains",[20]],["count",[10]],["count",[30]]]Expected: [null,true,true,false,true,2,true,false,2,1]

Input: [["RandomizedCollection",[]],["insert",[4]],["insert",[4]],["remove",[4]],["remove",[4]],["insert",[4]],["count",[4]]]Expected: [null,true,false,true,true,true,1]

Need a hint? (5 available)

The Affirm angle

This exact shape, put or get or delete plus a random access, is the single most repeated coding prompt in Affirm's loop, and the duplicate-and-weighted version is the hardest common variant. Two things win the round. First, state the swap-to-tail plan before you code it, and name why a middle delete would be O(n). Second, when you reach get_random, say out loud that duplicates make the uniform pick weighted by frequency for free, since that is the twist Affirm is really probing. Talk through the structure before you type, because the prompt is deliberately vague and the interviewer is grading whether you pinned it down.

Take it further

A common follow-up swaps frequency weighting for arbitrary weights: each value carries a weight that is not just its count, and get_random must respect it. The O(1) array pick no longer applies, so you keep a running prefix sum of weights and binary-search a random point into it, which lands each pick in O(log n). Another follow-up asks for deleteRandom, evicting a random element in O(1). That one is easy from here: pick a uniform slot, then run the same swap-to-tail removal you already wrote.

Show solution & walkthrough