Muestreo por reservorio (reservoir sampling)

De qué trata este capítulo

Todas las estructuras de datos de este libro asumían hasta ahora que podías tener tus datos en memoria — o al menos que sabías cuántos eran. Este capítulo final ataca el caso en que no puedes ni sabes: un stream de elementos de longitud desconocida, posiblemente infinita, que ves una sola vez, en orden, demasiado grande para guardarlo. ¿Cómo sacas de ahí una muestra aleatoria uniforme — con todos los elementos igual de probables — cuando nunca puedes verlo completo y no sabes qué tan largo es? El muestreo por reservorio es la respuesta elegante: una sola pasada, memoria fija (apenas la necesaria para guardar las k muestras) y una garantía demostrable de que, cuando el stream por fin termine, cada uno de sus n elementos — incluido el primero, que pasó hace miles de millones de elementos — tuvo exactamente la misma probabilidad k/n de quedar en tu muestra. En este capítulo lo construimos, vemos los elementos desfilar frente a un reservorio pequeño que los acepta o rechaza con una probabilidad que se encoge, y comprobamos que su uniformidad es real donde las alternativas obvias están muy sesgadas. Es un algoritmo pequeño y perfecto, y un cierre a la altura de un nivel — y de un libro — que trata sobre cómo la idea correcta vuelve rutinario lo que parecía imposible.

Un poco de historia

La forma más simple del muestreo por reservorio, el "Algoritmo R", la popularizó Jeffrey Vitter en un artículo de 1985, "Random Sampling with a Reservoir", aunque la idea de fondo ya circulaba antes (a Alan Waterman se le suele dar el crédito, vía The Art of Computer Programming de Knuth, donde aparece como ejercicio). La aportación de Vitter fueron tanto el análisis limpio como variantes más rápidas (el Algoritmo L y otros que saltan por encima de los elementos rechazados, convirtiendo las pasadas O(n) en O(k log(n/k)) al calcular de golpe cuántos elementos brincarse). La importancia del algoritmo creció muchísimo con el big data y los sistemas de streaming: cuando los datasets rebasaron la memoria y los datos empezaron a llegar como streams sin fin — logs web, clickstreams, telemetría de sensores, paquetes de red — el muestreo por reservorio se volvió la forma estándar de tomar una muestra justa sin tener que guardarlo todo. Viene integrado en frameworks de big data (Spark, Flink), en motores de bases de datos (para procesamiento aproximado de consultas y estadísticas) y en sistemas de monitoreo. Su atractivo duradero es que toma un problema que suena genuinamente difícil — muestrear de forma justa algo que no puedes guardar ni medir — y lo reduce a cinco líneas de código y una probabilidad exacta, que es justo por lo que es favorito tanto en la práctica como en entrevistas.

La intuición

Los intentos ingenuos fallan todos, y ver por qué afila la idea. ¿Quedarte con los primeros k elementos? Sesgado hacia el principio — los elementos posteriores nunca tienen oportunidad. ¿Elegir k índices aleatorios de antemano? No puedes, porque no conoces n. ¿Guardar todo y muestrear al final? Eso necesita memoria O(n) que no tienes, y con un stream infinito nunca termina. El requisito es estricto: una pasada, memoria O(k) y uniformidad — cada elemento, sin importar cuándo llegó, igual de probable.

El muestreo por reservorio lo cumple con una regla que parece casi demasiado simple. Guarda los primeros k elementos en un "reservorio". Después, para cada elemento siguiente (el i-ésimo, contando desde 1), consérvalo con probabilidad k/i; y cuando lo conserves, haz que reemplace a uno de los k elementos que ya están en el reservorio, elegido uniformemente al azar. Fíjate en que la probabilidad de aceptación se encoge conforme crece el stream: el elemento 100 se conserva con probabilidad k/100, el millonésimo con probabilidad k/1,000,000 — los elementos tardíos son individualmente menos probables de ser aceptados, que es exactamente lo que compensa el hecho de que los elementos tempranos han tenido más oportunidades de ser reemplazados. Lo mágico es que estos dos efectos se cancelan perfectamente: en todo momento, y al final de todo, cada elemento visto hasta ahí está en el reservorio con probabilidad exactamente k/(cuenta actual). Un elemento que entró temprano sobrevivió muchas rondas de posible reemplazo; un elemento que entró tarde tenía menos probabilidad de ser aceptado siquiera; y la aritmética cuadra de modo que todos, del primero al último, terminan con probabilidad k/n. El duelo lo confirma: la probabilidad de selección del muestreo por reservorio es completamente plana a lo largo de todas las posiciones del stream, mientras que "quedarse con los primeros k" solo selecciona el inicio y "aceptar con un ½ constante" sobre-selecciona el final.

Complejidad: cómo escala

El muestreo por reservorio es una sola pasada O(n) sobre el stream — un número aleatorio y una comparación por elemento — usando memoria O(k), independiente de la longitud n del stream (que nunca necesita conocer). El Algoritmo L de Vitter mejora el tiempo a O(k log(n/k)) calculando, después de cada aceptación, cuántos elementos saltarse antes de la siguiente aceptación (sacando ese hueco de la distribución correcta), evitando un volado por cada elemento rechazado — útil cuando las aceptaciones son raras en un stream larguísimo. Pero la propiedad estelar no es la velocidad, es la corrección bajo restricciones: uniformidad sin guardar nada. El duelo mide esa uniformidad de forma directa — la probabilidad empírica de que cada posición del stream termine en la muestra, para el muestreo por reservorio contra dos samplers de stream plausibles pero equivocados:

La línea del muestreo por reservorio es plana — cada posición se queda en k/n = 0.2, y la diferencia entre la posición más y la menos probable es de apenas 0.011, ruido estadístico. Esa planitud es la uniformidad: ninguna posición está favorecida, exactamente como se pedía. Los dos samplers sesgados muestran qué pasa cuando la probabilidad está mal. "Quedarse con los primeros k" es una función escalón — probabilidad 1 para las primeras seis posiciones, 0 para todas las demás, el sesgo más extremo posible (diferencia 1.0). "Aceptar con un ½ constante" (en lugar del k/i decreciente) sube de forma sostenida hacia el final, porque una tasa de aceptación fija deja que los elementos tardíos sigan sacando a los tempranos, sobre-representando la cola (diferencia 0.43). Ambos son samplers de stream que se ven razonables y están silenciosa y gravemente mal — que es justo el punto de la regla exacta k/i. La diferencia entre una muestra justa y una sesgada es una sola probabilidad bien elegida, y solo el muestreo por reservorio la acierta.

En qué es bueno y en qué no

El muestreo por reservorio es la herramienta correcta siempre que tengas que muestrear datos que no puedes guardar completos o cuyo tamaño no conoces de antemano: datos en streaming (clickstreams, telemetría, tráfico de red, líneas de log), archivos o tablas de base de datos demasiado grandes para la memoria (un solo barrido los muestrea de forma justa) y escenarios online donde debes mantener una muestra representativa conforme llegan los datos y estar listo para reportarla en cualquier momento. Se usa en pruebas A/B y muestreo de experimentos, en procesamiento aproximado de consultas en bases de datos, en muestreo de carga y rendimiento en sistemas de monitoreo, y en cualquier lugar donde haga falta una muestra uniforme de un dataset grande o sin límite en una sola pasada. Sus garantías son exactas (uniformidad real, no aproximada), su memoria es mínima y fija, y no necesita saber nada sobre la longitud del stream.

Donde es la herramienta equivocada es cuando sí puedes tener todos los datos y conoces su tamaño — ahí random.sample es más simple y hace lo mismo sin toda la maquinaria de streaming. Da una muestra uniforme sin pesos por defecto; el muestreo ponderado (elementos con probabilidades distintas) necesita la variante ponderada (A-Res/A-ExpJ), que es más elaborada. Muestrea sin reemplazo dentro del reservorio; muestrear con reemplazo, o hacer muestreo estratificado o por grupos, requiere adaptaciones. Y aunque es exactamente uniforme, una sola muestra aleatoria sigue siendo solo eso, una muestra — tiene varianza de muestreo, así que para un dataset fijo donde necesitas reproducibilidad o un subconjunto específico, otros métodos pueden quedar mejor. El nicho del muestreo por reservorio es preciso e importante: muestreo justo, de memoria fija y una sola pasada sobre streams y datos que no caben — el caso que nada más resuelve con tanta limpieza.

Los datos, o las entradas

El duelo mide empíricamente la probabilidad de que cada posición del stream termine en la muestra — a lo largo de decenas de miles de pruebas — para el muestreo por reservorio contra dos samplers de stream sesgados (quedarse con los primeros k y aceptación con ½ constante), mostrando la uniformidad plana del reservorio frente a su sesgo. La corrección se verifica de la misma forma con una tolerancia más fina: en 40,000 pruebas, la frecuencia empírica de selección de cada elemento debe quedar dentro de 0.02 del k/n teórico, la muestra siempre debe tener exactamente k elementos distintos del stream, y un stream más corto que k debe devolver todo. La animación corre un reservorio de tamaño 3 sobre un stream de diez elementos, mostrando cómo los primeros tres lo llenan y cómo cada elemento posterior es aceptado (reemplazando una posición aleatoria) o rechazado con la probabilidad decreciente k/i.

Constrúyelo, una función a la vez

Muestreo por reservorio — el algoritmo completo en una sola pasada:

def reservoir_sample(stream, k, rng=random):
    """A uniform random k-sample of `stream` in one pass, O(k) memory, without knowing the stream's
    length. Fill the reservoir with the first k items. For each later item at index i (0-based), pick
    a random slot j in [0, i]; if j lands inside the reservoir (j < k), the new item replaces slot j —
    which happens with probability k/(i+1), exactly the acceptance rate that keeps the sample uniform.
    Returns the reservoir."""
    reservoir = []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)                # the first k items fill the reservoir
        else:
            j = rng.randint(0, i)                 # uniform in [0, i]; accepts with prob k/(i+1)
            if j < k:
                reservoir[j] = item               # replace a uniformly random held item
    return reservoir


def reservoir_trace(stream, k, rng=random):
    """Same algorithm, but record each step (accepted?, which slot, the reservoir after) — for the
    animation. Returns (reservoir, steps)."""
    reservoir, steps = [], []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)
            steps.append({"i": i, "item": item, "action": "fill", "slot": i, "prob": 1.0,
                          "reservoir": reservoir[:]})
        else:
            j = rng.randint(0, i)
            prob = k / (i + 1)
            if j < k:
                reservoir[j] = item
                steps.append({"i": i, "item": item, "action": "accept", "slot": j, "prob": prob,
                              "reservoir": reservoir[:]})
            else:
                steps.append({"i": i, "item": item, "action": "reject", "slot": None, "prob": prob,
                              "reservoir": reservoir[:]})
    return reservoir, steps

Míralo funcionar

Aquí tienes un reservorio de tamaño 3 muestreando un stream de diez elementos, de la A a la J. La celda de más a la izquierda (in) es el elemento que está llegando; las tres celdas a su derecha son el reservorio. Los primeros tres elementos — A, B, C — simplemente llenan el reservorio vacío. Y luego viene lo interesante: cada elemento nuevo se conserva con probabilidad k/i (se ve en el pie de figura, y va encogiéndose conforme crece el stream — 3/4 para el 4º elemento, 3/5 para el 5º, y así). Cuando un elemento es aceptado (verde), reemplaza a uno de los tres elementos guardados elegido uniformemente al azar, sacándolo; cuando es rechazado (rojo), el reservorio queda igual. Fíjate en cómo los elementos tardíos son individualmente menos probables de entrar — pero si entran, pueden desalojar a un sobreviviente temprano. Cuando el stream termina, los tres elementos que quedan son una muestra genuinamente uniforme: cada uno de los diez tenía probabilidad 3/10 de estar entre ellos, sin importar qué tan temprano o tarde llegó, que es justo lo que la aceptación decreciente k/i está diseñada para garantizar:

El código completo

La pestaña de implementación propia es el muestreo por reservorio (y la versión con trazas para la animación); la pestaña de librería es la referencia de guardar-todo-y-luego-random.sample (exacta pero con memoria O(n) y necesita todos los datos) más los dos samplers sesgados del duelo. Cambia entre ellas — random.sample es lo que usarías cuando sí puedes tener los datos en memoria; el muestreo por reservorio es la respuesta de cinco líneas para cuando no puedes, y compararlos muestra exactamente alrededor de qué restricción está diseñada la versión con reservorio.

"""Reservoir sampling — draw a uniform random sample of k items from a stream of UNKNOWN, possibly
infinite length, in a single pass, using only enough memory to hold the k samples. It answers a
question the earlier structures can't touch: how do you sample fairly from data too big to store, or
still arriving, when you can't hold it all and don't even know how much there is? Log files, sensor
streams, click events, database scans too large for memory — all of these you see once, in order, and
must sample without a second pass.

The naive approaches fail: keeping the first k biases toward the beginning; you can't pick k random
INDICES because you don't know n; storing everything to sample at the end needs O(n) memory you don't
have. Reservoir sampling (Algorithm R) solves it with a beautifully simple rule. Keep the first k
items. Then for the i-th item (counting from 1), keep it with probability k/i, and if you keep it,
have it REPLACE a uniformly random one of the k currently held. That's it — one pass, O(k) memory, and
a provable guarantee that at the end EVERY item the stream produced, no matter how long ago, has
exactly k/n probability of being in the sample.
"""
import random


# region: reservoir
def reservoir_sample(stream, k, rng=random):
    """A uniform random k-sample of `stream` in one pass, O(k) memory, without knowing the stream's
    length. Fill the reservoir with the first k items. For each later item at index i (0-based), pick
    a random slot j in [0, i]; if j lands inside the reservoir (j < k), the new item replaces slot j —
    which happens with probability k/(i+1), exactly the acceptance rate that keeps the sample uniform.
    Returns the reservoir."""
    reservoir = []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)                # the first k items fill the reservoir
        else:
            j = rng.randint(0, i)                 # uniform in [0, i]; accepts with prob k/(i+1)
            if j < k:
                reservoir[j] = item               # replace a uniformly random held item
    return reservoir


def reservoir_trace(stream, k, rng=random):
    """Same algorithm, but record each step (accepted?, which slot, the reservoir after) — for the
    animation. Returns (reservoir, steps)."""
    reservoir, steps = [], []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)
            steps.append({"i": i, "item": item, "action": "fill", "slot": i, "prob": 1.0,
                          "reservoir": reservoir[:]})
        else:
            j = rng.randint(0, i)
            prob = k / (i + 1)
            if j < k:
                reservoir[j] = item
                steps.append({"i": i, "item": item, "action": "accept", "slot": j, "prob": prob,
                              "reservoir": reservoir[:]})
            else:
                steps.append({"i": i, "item": item, "action": "reject", "slot": None, "prob": prob,
                              "reservoir": reservoir[:]})
    return reservoir, steps
# endregion

"""The reference and the contrasts. Three ways to think about sampling k from a stream:

- `store_all_sample` is what you'd do if you COULD hold everything — buffer the whole stream, then
  `random.sample`. It's exact and uniform (the correctness reference for reservoir sampling's
  uniformity), but it needs O(n) memory and a known, finite length — exactly what reservoir sampling
  avoids.
- `keep_first_k` and `constant_half` are BIASED stream samplers, to show why the k/i acceptance
  probability is not arbitrary: keeping the first k over-weights the start, and accepting every item
  with a constant 1/2 over-weights the end. Comparing their per-position selection frequency against
  reservoir sampling's flat line is the face-off.

In production you'd use reservoir sampling directly (it's short) or a streaming framework's built-in;
`random.sample` is the reference for finite in-memory data.
"""
import random


# region: reference
def store_all_sample(stream, k):
    """Buffer the entire stream, then random.sample — exact and uniform, but O(n) memory and needs the
    full data in hand. The reference reservoir sampling matches without either requirement."""
    data = list(stream)
    if k >= len(data):
        return data[:]
    return random.sample(data, k)
# endregion


# region: biased
def keep_first_k(stream, k):
    """A biased sampler: just keep the first k items. Positions 0..k-1 are always selected, everything
    after is never selected — maximally biased toward the start."""
    out = []
    for item in stream:
        if len(out) < k:
            out.append(item)
        else:
            break
    return out


def constant_half(stream, k, rng=random):
    """A biased sampler: fill k, then accept every later item with a CONSTANT probability 1/2 (not
    k/i), replacing a random slot. Because the acceptance rate doesn't shrink as the stream grows,
    later items are over-represented — the sample skews toward the end."""
    reservoir = []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)
        elif rng.random() < 0.5:
            reservoir[rng.randrange(k)] = item
    return reservoir
# endregion

Implementación propia vs librería

El muestreo por reservorio es una nota de cierre perfecta porque destila en cinco líneas toda la lección del nivel avanzado: la idea correcta vuelve rutinario lo imposible. Muestrear de forma justa un stream infinito que no puedes guardar suena a que no debería ser posible — y los intentos ingenuos fallan todos — y sin embargo una sola probabilidad exacta, k/i, hace que salga limpio, con una demostración que telescopa a k/n. Ese es el mismo espíritu del filtro de Bloom (representar un conjunto sin guardarlo), la skip list (balance sin rebalanceo) y el k-d tree (buscar en el espacio sin recorrerlo todo): una estructura ingeniosa o una probabilidad ingeniosa te compran algo que la fuerza bruta no puede. Y el duelo deja clarísimo un tema de todo el libro — que un algoritmo que se ve plausible puede estar silenciosa y gravemente mal (el muestreo con ½ constante se ve bien y se sesga feo), así que la corrección hay que razonarla y verificarla, no darla por hecha. En producción escribirías el muestreo por reservorio inline (es más corto que importar cualquier cosa) o usarías el que ya trae tu framework de streaming; entenderlo es entender cómo los sistemas de big data muestrean lo ilimitado, y por qué la probabilidad exacta, y no una que parezca razonable, es lo que hace justa a una muestra.

Dónde te lo vas a encontrar

El muestreo por reservorio corre por toda la infraestructura de big data y streaming. Los frameworks de procesamiento de datos (el sample de Spark, Flink) lo usan para muestrear datasets enormes o en streaming en una sola pasada. Las bases de datos lo usan para procesamiento aproximado de consultas, recolección de estadísticas y muestreo de tablas grandes que no caben en memoria. Los sistemas de monitoreo y observabilidad (métricas, tracing, muestreo de logs) lo usan para mantener una muestra representativa de eventos de alto volumen bajo memoria fija — los sistemas de tracing distribuido muestrean spans así. Las plataformas de pruebas A/B y experimentación muestrean usuarios o eventos con él. Los pipelines de machine learning lo usan para submuestrear datos de entrenamiento enormes o en streaming. Los sistemas de red muestrean paquetes con él para análisis. Y es un problema clásico de entrevistas y de programación competitiva precisamente porque prueba si entiendes la probabilidad exacta. En cualquier lugar donde haga falta una muestra justa de un stream o de un dataset demasiado grande en una sola pasada, el muestreo por reservorio — o sus variantes ponderadas y optimizadas con saltos — muy probablemente está haciendo el trabajo.

Conclusiones

El muestreo por reservorio saca una muestra aleatoria uniforme de tamaño k de un stream de longitud desconocida en una sola pasada y con memoria O(k): guarda los primeros k, y luego conserva el i-ésimo elemento con probabilidad k/i, reemplazando a un elemento guardado elegido uniformemente al azar. La probabilidad de aceptación decreciente compensa exactamente la mayor exposición al reemplazo de los elementos tempranos, así que cada elemento — del primero al último — termina con probabilidad k/n, demostrable con un producto telescópico. No necesita ni la longitud del stream ni todos los datos, y su uniformidad es exacta donde las alternativas plausibles (quedarse con los primeros k, aceptación constante) están muy sesgadas. Es la respuesta limpia al muestreo justo de lo que no se puede guardar.

Con el muestreo por reservorio cierra el libro. Cincuenta y seis capítulos han construido el núcleo de las estructuras de datos y los algoritmos desde cero — desde cómo medimos el costo, pasando por las estructuras lineales, el hashing, la recursión y el ordenamiento, los árboles, los grafos, las cadenas, los paradigmas de diseño de algoritmos, y finalmente estas estructuras avanzadas y probabilísticas — cada una construida a mano, vista operando sobre datos reales y medida contra la librería con la que compite. El hilo conductor ha sido que la estructura correcta o la idea correcta es lo que separa lo tratable de lo intratable, lo justo de lo sesgado, lo rápido de lo lento: una tabla en lugar de recalcular, un carril express en lugar de caminar, una heurística en lugar de una inundación, una probabilidad en lugar de una garantía. Ese criterio — saber qué idea pide un problema, y por qué — es para lo que ha servido todo el libro.