Коротко

Reservoir sampling, или резервуарная выборка, — короткий алгоритм с красивым доказательством. Его любят на собеседованиях по двум причинам: код занимает пять строк, а объяснить, почему он работает, может не каждый. Именно объяснение и проверяют.

Постановка задачи

Есть поток элементов: строки огромного лога, записи из сети, узлы связного списка. Мы видим элементы по одному и не знаем, сколько их будет. Нужно в конце вернуть один элемент так, чтобы каждый был выбран с одинаковой вероятностью.

Ограничения, которые делают задачу интересной:

• поток проходим один раз, вернуться назад нельзя;
• сохранить всё нельзя: данных может быть больше, чем памяти;
• длина $n$ становится известна только в конце.

Если бы $n$ было известно, решение тривиально: сгенерировать индекс random.randint(0, n - 1) и дойти до него. Если бы можно было сохранить всё — сложить в массив и выбрать. Резервуарная выборка нужна, когда нельзя ни то, ни другое.

Наивные решения и почему они не подходят

Два прохода. Первый проход считает длину, второй идёт до случайного индекса. Для связного списка это работает, но для потока — нет: второго прохода не будет.

Сохранить всё. Память $O(n)$. Для списка из $10^4$ узлов не страшно, для терабайтного лога — невозможно.

Брать каждый элемент с вероятностью 1/2. Тогда последний элемент выигрывает с вероятностью 1/2, а первый — с $1/2^{n-1}$. Распределение неравномерное.

Нужна вероятность, которая зависит от номера элемента. Ровно это и делает алгоритм R.

Алгоритм R для одного элемента

Идём по потоку и нумеруем элементы с единицы. Встретив $i$-й элемент, с вероятностью $1/i$ делаем его текущим ответом, иначе оставляем прежний.

import random


def reservoir_one(stream):
    chosen = None
    for i, x in enumerate(stream, start=1):
        if random.randint(1, i) == 1:  # вероятность 1/i
            chosen = x
    return chosen

Первый элемент берётся всегда ($1/1$). Второй заменяет его с вероятностью $1/2$. Третий — с вероятностью $1/3$. И так далее.

Проверим эмпирически: 100 000 запусков на потоке range(5).

from collections import Counter

random.seed(1)
print(sorted(Counter(reservoir_one(range(5)) for _ in range(100_000)).items()))
# [(0, 20181), (1, 19925), (2, 19870), (3, 19901), (4, 20123)]

Каждое число выпало около 20 000 раз, то есть с частотой $1/5$. Теперь докажем, что это не совпадение.

Доказательство равновероятности

Посчитаем вероятность, что в конце выбран элемент с номером $i$. Для этого должны произойти два события:

1. На шаге $i$ элемент был взят: вероятность $\frac{1}{i}$.
2. На каждом следующем шаге $j = i+1, \dots, n$ его не заменили: вероятность $1 - \frac{1}{j} = \frac{j-1}{j}$.

Решения на разных шагах независимы, поэтому вероятности перемножаются:

Произведение телескопическое: числитель каждой дроби сокращается со знаменателем предыдущей, и остаётся $1/n$. Это верно для любого $i$, значит, все элементы равновероятны.

Доказательство по индукции

Интервьюеры иногда просят другую формулировку, через инвариант. Утверждение: после обработки $m$ элементов каждый из них является текущим ответом с вероятностью $1/m$.

• База, $m = 1$: первый элемент взят с вероятностью 1.
• Переход $m \to m+1$: новый элемент становится ответом с вероятностью $\frac{1}{m+1}$. Любой из старых был ответом с вероятностью $\frac{1}{m}$ и остался им с вероятностью $\frac{m}{m+1}$. Итого $\frac{1}{m} \cdot \frac{m}{m+1} = \frac{1}{m+1}$.

Инвариант выполняется на каждом шаге, поэтому поток можно оборвать в любой момент, и текущий ответ будет равномерно случайным среди уже прочитанного. Это важное свойство для бесконечных стримов: выборка корректна «здесь и сейчас».

Выборка k элементов

Обобщение: нужно выбрать $k$ различных элементов (без повторов) так, чтобы каждое подмножество размера $k$ было равновероятным. Резервуар — это массив на $k$ мест.

1. Первые $k$ элементов кладём в резервуар как есть.
2. Для элемента с индексом $i \ge k$ (нумерация с нуля) генерируем j от 0 до i включительно.
3. Если j < k, заменяем reservoir[j] новым элементом.

def reservoir_k(stream, k):
    reservoir = []
    for i, x in enumerate(stream):
        if i < k:
            reservoir.append(x)
        else:
            j = random.randint(0, i)  # 0..i включительно, всего i+1 вариантов
            if j < k:
                reservoir[j] = x
    return reservoir

Новый элемент попадает в резервуар с вероятностью $\frac{k}{i+1}$, а выбрасывает случайный из $k$ старых, потому что j при условии j < k равномерно распределён на 0..k-1.

Почему вероятность равна k/n

Снова считаем для элемента с номером $t$ (с единицы, $t > k$). Он попал в резервуар с вероятностью $\frac{k}{t}$. На шаге $j > t$ его вытесняют, если новый элемент принят ($\frac{k}{j}$) и выбрал именно его место ($\frac{1}{k}$): вероятность $\frac{1}{j}$. Значит, он выживает с вероятностью $\frac{j-1}{j}$, и снова получается телескоп:

Для первых $k$ элементов всё то же самое, только первый множитель равен 1, а произведение идёт от $j = k+1$: $\frac{k}{k+1} \cdots \frac{n-1}{n} = \frac{k}{n}$.

Эксперимент на range(6) при $k = 2$: каждое число должно попасть в выборку с частотой $2/6 = 1/3$.

random.seed(1)
c = Counter()
for _ in range(100_000):
    c.update(reservoir_k(range(6), 2))
print(sorted(c.items()))

каждое число около 33 300 раз из 100 000

Можно проверить и сильное утверждение, что равновероятны именно пары. Для `range(4)` и $k = 2$ все шесть пар выпадают примерно по 16 700 раз из 100 000, то есть по $1/6$.

### Что если k больше длины потока

Тогда резервуар просто не заполнится, и функция вернёт все элементы. Обычно это правильное поведение, но в интервью стоит проговорить его вслух.

## Типичные ошибки

### Ошибка на единицу в диапазоне

Самая частая и самая коварная ошибка — написать `random.randint(0, i - 1)` вместо `random.randint(0, i)`. Код выглядит правдоподобно и не падает, но распределение ломается. Тот же эксперимент на `range(6)`, $k = 2$:

```text
правильно:   0: 33 267  1: 33 225  2: 33 217  3: 33 315  4: 33 628  5: 33 348
ошибка:      0: 20 076  1: 19 954  2: 39 671  3: 40 227  4: 40 235  5: 39 837

На элементе с индексом 2 «ошибочный» код выбирает j из {0, 1}, условие j < 2 выполняется всегда, и элемент попадает в резервуар со стопроцентной вероятностью. Ранние элементы вытесняются слишком часто.

Помните: в Python random.randint(a, b) включает обе границы, а random.randrange(a, b) — нет. В Java rand.nextInt(n) даёт 0..n-1, в JavaScript Math.floor(Math.random() * n) — тоже 0..n-1. Для элемента с индексом i с нуля нужно $i + 1$ вариантов.

Нумерация с нуля в версии для одного элемента

if random.randint(0, i) == 0 при enumerate(stream) с нуля — правильно. if random.randint(1, i) == 1 при нумерации с нуля упадёт на первом же элементе: randint(1, 0) бросает исключение. Выберите одну нумерацию и держитесь её.

Проверка на «выглядит случайно»

Запустить функцию три раза и посмотреть на разные ответы — не проверка. Нужен счётчик частот на десятках тысяч запусков, как выше. На собеседовании достаточно сказать, что вы бы так проверили.

Задачи с LeetCode

Linked List Random Node (LeetCode 382)

Дан односвязный список. Метод getRandom() должен возвращать значение случайного узла, каждый узел с вероятностью $1/n$. Follow-up из условия: что если список огромный, длина неизвестна, а дополнительной памяти нет?

Ответ на follow-up — резервуар из одного элемента, где поток — это обход списка.

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


class Solution:
    def __init__(self, head):
        self.head = head

    def getRandom(self):
        result, node, i = None, self.head, 0
        while node:
            i += 1
            if random.randint(1, i) == 1:
                result = node.val
            node = node.next
        return result


s = Solution(ListNode(1, ListNode(2, ListNode(3))))
print(s.getRandom())  # 1, 2 или 3, каждое с вероятностью 1/3

Каждый вызов — $O(n)$ времени и $O(1)$ памяти. Альтернатива — при создании скопировать значения в массив и отвечать за $O(1)$, тратя $O(n)$ памяти. Стоит назвать оба варианта и сказать, когда какой лучше: много вызовов на небольшом списке — массив; один-два вызова на гигантском — резервуар.

Random Pick Index (LeetCode 398)

Дан массив с повторами. Метод pick(target) возвращает случайный индекс, где стоит target, каждый такой индекс равновероятно.

Здесь «поток» — это только элементы, равные target. Счётчик увеличиваем лишь на них.

class Solution:
    def __init__(self, nums):
        self.nums = nums

    def pick(self, target):
        result, count = -1, 0
        for i, x in enumerate(self.nums):
            if x == target:
                count += 1
                if random.randint(1, count) == 1:
                    result = i
        return result


s = Solution([1, 2, 3, 3, 3])
print(s.pick(3))  # 2, 3 или 4, каждое с вероятностью 1/3
print(s.pick(1))  # 0

Частая ошибка: использовать в вероятности индекс массива i, а не число найденных совпадений count. Тогда вероятность зависит от позиции в массиве, и равномерность пропадает.

Решение через хеш-таблицу «значение → список индексов» отвечает за $O(1)$, но хранит $O(n)$ индексов. На LeetCode оно проходит, однако задача задумана именно под резервуар, и интервьюер почти наверняка спросит про вариант без дополнительной памяти.

Взвешенная выборка: A-Res

Иногда элементы неравноценны: у каждого есть вес $w_i > 0$, и вероятность выбора должна от него зависеть. Например, выбрать запросы для отладки пропорционально их длительности.

Алгоритм A-Res (Efraimidis и Spirakis, 2006) делает это одним проходом:

1. Для каждого элемента генерируем $u$ равномерно на $(0, 1)$ и считаем ключ $u^{1/w}$.
2. Оставляем $k$ элементов с наибольшими ключами.

Хранить «$k$ наибольших» удобно в куче с минимумом наверху: новый элемент вытесняет корень, если его ключ больше. Это та же схема, что в задаче Top K Frequent Elements.

import heapq


def weighted_reservoir(stream, k):
    heap = []  # (ключ, элемент), минимальный ключ наверху
    for x, w in stream:
        if w <= 0:
            continue
        key = random.random() ** (1 / w)
        if len(heap) < k:
            heapq.heappush(heap, (key, x))
        elif key > heap[0][0]:
            heapq.heapreplace(heap, (key, x))
    return [x for _, x in heap]


random.seed(1)
items = [("a", 1), ("b", 2), ("c", 7)]
print(sorted(Counter(weighted_reservoir(items, 1)[0] for _ in range(100_000)).items()))
# [('a', 9966), ('b', 19996), ('c', 70038)]

При $k = 1$ вероятность выбора ровно $w_i / \sum w$: веса 1, 2, 7 дали частоты 10%, 20%, 70%. При $k > 1$ это выборка без возвращения: сначала с такой вероятностью выбирается первый элемент, потом среди оставшихся — второй, и итоговая вероятность попасть в выборку уже не строго пропорциональна весу. Об этом стоит сказать, если интервьюер спросит про точность.

Время $O(n \log k)$, память $O(k)$. Почему ключ именно $u^{1/w}$: чем больше вес, тем ближе $1/w$ к нулю и тем ближе ключ к единице, то есть у тяжёлого элемента больше шансов оказаться наверху. Строгое доказательство есть в оригинальной статье, на собеседовании его обычно не требуют.

Как ускорить: идея алгоритма L

Алгоритм R вызывает генератор случайных чисел на каждом элементе. Когда $n$ намного больше $k$, почти все вызовы заканчиваются отказом. Алгоритм L (Li, 1994) вместо этого сразу вычисляет, сколько элементов пропустить до следующей замены, и генерирует случайное число только на заменах. Ожидаемое число таких событий $O(k(1 + \log(n/k)))$.

Для собеседования достаточно знать, что такая оптимизация существует, и уметь объяснить, откуда берётся выигрыш. Писать её по памяти не просят.

Где резервуарная выборка встречается в жизни

• Логи и мониторинг. Сохранить 1000 случайных запросов за сутки из миллиардов, не зная заранее их количества.
• Базы данных. Оценка статистики по таблице, TABLESAMPLE и сбор выборки для планировщика запросов.
• Машинное обучение. Буфер воспроизведения в обучении с подкреплением и потоковые датасеты, которые не помещаются в память.
• Распределённые системы. Каждый узел ведёт свой резервуар, а потом выборки объединяют с учётом размеров потоков.

Рядом стоят другие потоковые алгоритмы с похожей философией «приближённо, но в постоянной памяти»: фильтр Блума для проверки принадлежности и голосование Бойера — Мура для поиска элемента большинства за один проход.

Как отвечать на собеседовании

1. Уточните условия: известна ли длина, можно ли хранить данные, сколько будет вызовов.
2. Назовите тривиальные решения (два прохода, массив) и их цену.
3. Предложите алгоритм R: «$i$-й элемент беру с вероятностью $1/i$».
4. Докажите через телескопическое произведение или по индукции. Этот шаг важнее кода.
5. Обобщите на $k$ элементов и назовите сложность: $O(n)$ времени, $O(k)$ памяти.
6. Скажите, как проверили бы: счётчик частот на большом числе запусков.

Если после этого спросят про веса, упомяните A-Res с ключом $u^{1/w}$ и кучей на $k$ элементов. Для смежной темы о случайности в алгоритмах посмотрите Quickselect: там случайный выбор опорного элемента даёт ожидаемое линейное время.