398. Random Pick Index

398

MediumRandomization, Math and Miscellaneous (FAANG add-on)Single-pass uniform sampling over a streamReservoir sampling (k=1)

Random Pick Index

Given an integer array nums that may contain duplicates, implement pick(target): return a random index i such that nums[i] == target, with every matching index equally likely to be returned. It is guaranteed target exists in nums.

Open official problem prompt ↗

In plain English

Return one index uniformly at random among all positions equal to target, using only constant extra memory per query.

Picture it like this

Interviewing candidates one at a time and, at the c-th qualified candidate, giving them the job with probability 1/c — everyone who ever qualified ends up equally likely to hold the offer.

Example

Input: nums = [1,2,3,3,3]; calls = pick(3), pick(1), pick(3)
Output: [4, 0, 2]
Why: pick(3) may return any of indices 2, 3, or 4 each with probability 1/3 (here 4); pick(1) must return the only match, index 0; pick(3) again returns one of 2,3,4 (here 2).

Constraints

1 <= nums.length <= 2 * 10^4-2^31 <= nums[i] <= 2^31 - 1target is an integer that exists in numsAt most 10^4 calls to pick

Pattern lesson

See the pattern, then code

Reservoir sampling (k=1)

↗

Randomization, Math and Miscellaneous (FAANG add-on)

Uniform random selection, sampling streams, number theory, or precision-heavy parsing.. Scan the array counting matches; when you meet the c-th match, keep it as the answer with probability 1/c. This overwrites earlier picks just often enough that after the pass each of the c matches survives with probability exactly 1/c.

New words, made simpleKnow these before the algorithm

Reservoir sampling: A family of algorithms for sampling k items uniformly from a stream of unknown or large length in one pass.
Reservoir of size 1: The special case where you keep a single winner and replace it with probability 1/count.
Uniform over matches: Every index whose value equals target has identical probability of being returned.

Approaches from first idea to best ideaCompare the trade-offs

ApproachHow it thinksCost

Collect all indices then pick

Simple and correct; uses extra memory proportional to matches, which the follow-up asks to avoid.

Build a list of every index equal to target and return random.choice.

Time O(n) per pickSpace O(m) for the m matches (or O(n) precomputed)

Precompute value -> list of indices

Fastest queries, but O(n) memory and heavy preprocessing.

Hash map from value to its index list built once in the constructor.

Time O(1) expected per pickSpace O(n)

Best choiceReservoir sampling

Meets the constant-space constraint with a clean uniform guarantee.

One pass counting matches, keep index i with probability 1/count.

Time O(n) per pickSpace O(1)

The rule we keep true

Invariant

After examining the c-th match, the stored result holds each of those c matching indices with probability exactly 1/c.

Why this is correct

Reasoning

The c-th match is kept with probability 1/c. An earlier match j (with j < c) survives only if it was chosen when seen (prob 1/j) and then not replaced by matches j+1..c: (1/j)·(j/(j+1))·((j+1)/(j+2))···((c-1)/c) telescopes to 1/c. So all c matches share probability 1/c.

The algorithm in three movesSay these aloud before coding

1Iterate through nums tracking how many matches seen so far (count)

match@2: count=1, keep prob 1/1 -> res=2

2On each match increment count

match@3: count=2, keep prob 1/2 -> maybe res=3

3Replace the stored answer with the current index with probability 1/count

match@4: count=3, keep prob 1/3 -> res=4

4After the full pass, return the stored index

Visual trace

Follow the data, one decision at a time

Input snapshotHighlighted cells are involved in this example

1 · Readtarget=3

2 · AskReset counters

3 · Update statecount=0, result=-1

4 · Resultready

Key takeaway

Each new match at count c replaces the answer with probability 1/c, keeping the choice uniform.

Code walkthrough

Read the solution in small chunks

Python 3

Do not memorize the whole program. Connect each group of lines to one job in the algorithm.

1
Lines 4-5Store reference
Keep the array; no preprocessing so memory stays O(1).
2
Lines 8-9Counters
count tracks matches seen; result holds the current winner.
3
Lines 10-14Reservoir step
randint(1,count)==1 is true with probability 1/count, replacing the winner uniformly.

Make it stick

Edge cases, mistakes, and one self-check

Edge cases to test

target appears exactly once — count stays 1 and that index is always returned
All elements equal target — reduces to uniform pick over the whole array
target at the very end still gets its fair 1/count chance

Common beginner mistakes

Using random.random() < 1.0/count without care for count=1 rounding — randint(1,count)==1 is cleaner
Resetting result to a stale value between calls — must reinitialize each pick
Assuming the first match should always win — it must be replaceable

Check your understanding

Why does keeping the c-th match with probability 1/c make every match equally likely, not just the last one?

Random Pick Index

Random Pick Index

See the pattern, then code

Invariant

Reasoning

Follow the data, one decision at a time

Read the solution in small chunks

Edge cases, mistakes, and one self-check

Edge cases to test

Common beginner mistakes

More in Randomization, Math and Miscellaneous (FAANG add-on)