Select a random number from stream

Problem

Given a stream of numbers, select a random number from the stream such that each number has an equal probability of being chosen. You are allowed to use only O(1) space, and the input is in the form of a stream, so you cannot store previously seen numbers.

Examples

Example 1

Input: Stream = [5, 7, 9]
Output: 5 (or 7 or 9, each with probability 1/3)
Explanation: Each number in the stream is equally likely to be chosen.

Example 2

Input: Stream = [10, 20, 30, 40]
Output: 20 (or 10 or 30 or 40, each with probability 1/4)
Explanation: Each number in the stream is equally likely to be chosen.

Solution

Method 1 – Reservoir Sampling

Intuition

The key idea is to maintain a single candidate answer as you process the stream. For each new number, you randomly decide whether to replace the current answer, ensuring that every number has an equal chance of being selected, regardless of the stream length.

Approach

1. Initialize ans to None.
2. For each number in the stream, increment a counter i.
3. Generate a random integer between 1 and i (inclusive).
4. If the random integer is 1, replace ans with the current number.
5. After the stream ends, ans is the randomly selected number.

Code

C++

#include <vector>
class Solution {
public:
  int selectRandom(const std::vector<int>& arr) {
    int ans = -1, i = 0;
    auto it = arr.begin();
    while (it != arr.end()) {
      ++i;
      if (rand() % i == 0) ans = *it;
      ++it;
    }
    return ans;
  }
};

Go

func SelectRandom(arr []int) int {
  ch := make(chan int)
  go func() {
    for _, v := range arr {
      ch <- v
    }
    close(ch)
  }()
  ans := -1
  i := 0
  for num := range ch {
    i++
    if rand.Intn(i) == 0 {
      ans = num
    }
  }
  return ans
}

Java

import java.util.Iterator;
import java.util.Arrays;

class Solution {
  public int selectRandom(int[] arr) {
    Iterator<Integer> stream = Arrays.stream(arr).iterator();
    int ans = -1, i = 0;
    while (stream.hasNext()) {
      int num = stream.next();
      i++;
      if (Math.random() < 1.0 / i) ans = num;
    }
    return ans;
  }
}

Kotlin

class Solution {
  fun selectRandom(arr: IntArray): Int {
    val stream = arr.asSequence()
    var ans = -1
    var i = 0
    for (num in stream) {
      i++
      if (Math.random() < 1.0 / i) ans = num
    }
    return ans
  }
}

Python

class Solution:
  def select_random(self, arr: list[int]) -> int:
    def stream():
      for v in arr:
        yield v
    ans = None
    for i, num in enumerate(stream(), 1):
      if random.randint(1, i) == 1:
        ans = num
    return ans

Rust

use rand::Rng;
impl Solution {
  pub fn select_random(arr: &[i32]) -> i32 {
    let mut ans = None;
    let mut i = 0;
    let mut stream = arr.iter();
    while let Some(&num) = stream.next() {
      i += 1;
      if rand::thread_rng().gen_range(1..=i) == 1 {
        ans = Some(num);
      }
    }
    ans.unwrap()
  }
}

TypeScript

class Solution {
  selectRandom(arr: number[]): number {
    function* stream() {
      for (const v of arr) yield v;
    }
    let ans: number | undefined = undefined;
    let i = 0;
    for (const num of stream()) {
      i++;
      if (Math.floor(Math.random() * i) === 0) ans = num;
    }
    return ans!;
  }
}

Complexity

• ⏰ Time complexity: O(n) — Each element in the stream is processed once.
• 🧺 Space complexity: O(1) — Only a constant number of variables are used.