Пример реализации reservoir sampling на Python
23 мая 2025
Большинство программистов не раз сталкивались с задачей, где нужно выбрать случайный элемент из массива. Может потребоваться выбрать не один элемент, а несколько, что не намного сложнее. Но что делать, если вместо массива на входе у нас поток данных неизвестной длины, а то и вовсе бесконечный?
Алгоритмы для решения подобных задач называются reservoir sampling, или резервуарная выборка. Как ни странно, здесь можно придумать больше одного алгоритма. Далее я опишу решение, которое на мой взгляд является самым простым и понятным.
Для начала рассмотрим упрощенную задачу. Есть поток чисел. Длина потока нам неизвестна, мы просто видим по одному числу за раз. Нужно выбрать случайное число, притом любое число должно возвращаться с равной вероятностью. Здесь предполагается, что все числа в потоке уникальные. Да и вообще, числа взяты для примера. С тем же успехом это могли бы быть яблоки или мячики.
Когда мы видим первое число, то запоминаем его. Если поток закончился, то просто возвращаем это число. Если на входе получили второе число, то с вероятностью 1/2 мы должны вернуть первое число, и с вероятность 1/2 — второе. Если получили третье, то каждое число имеет вероятность 1/3, и т.д. Размышляя таким образом, мы приходим к следующему алгоритму. Очередное число на входе замещает собой то, что мы держим в памяти, с вероятностью 1/S, где S — это количество видимых на данный момент чисел.
Немного поразмыслив, не трудно обобщить алгоритм для случая, когда нужно выбрать не одно число, а несколько. Число, которое мы запоминали в случае упрощенной задачи, можно рассматривать, как «окно» размером в один элемент. Нет причин, почему в окне не может храниться два, три, или какое-то W число элементов. Если окно еще не заполнено, то очередной элемент добавляется в него безусловно. Иначе элемент попадает в него с вероятностью W/S. При этом замещается случайный элемент в окне.
Можно показать аналитически, что любой элемент потока попадет в выборку с равной вероятностью. Заинтересованные читатели найдут соответствующие выкладки в приложении к посту.
Переведем описанный алгоритм на язык Python:
from random import random
class ReservoirSampler:
def __init__(self, nsamples):
self.nsamples = nsamples
self.nseen = 0
self.result = []
def sample(self, item):
self.nseen += 1
if self.nseen <= self.nsamples:
self.result += [item]
else:
if random() < self.nsamples/self.nseen:
i = int(random()*self.nsamples)
self.result[i] = item
max_i = 99
rcount = [ 0 for i in range(0, max_i+1) ]
for ntest in range(0, 10000):
sampler = ReservoirSampler(10)
for i in range(0, max_i+1):
sampler.sample(i)
for r in sampler.result:
rcount[r] += 1
for i in range(0, max_i+1):
print("{} -> {}".format(i, rcount[i]))
Помимо реализации класса ReservoirSampler здесь приводится код тестирования. В нем мы десять тысяч раз выбираем 10 случайных чисел из потока в 100 чисел и проверяем, как часто заданное число попало в выборку. Легко удостовериться, что каждое число попало в выборку около 1000 раз, или с вероятностью ~10%. Характерно, что осмысленный результат возвращается даже для nsamples ≤ 0.
Пусть реализация на Python и не выглядит особо захватывающей, зато она наглядно объясняет суть алгоритма и позволяет убедиться в его корректности. Reservoir sampling находит практическое применение при разработке СУБД, в обработке временных рядов и во многих других задачах.
Вы можете прислать свой комментарий мне на почту, или воспользоваться комментариями в Telegram-группе.