Logarithmic entropy usage for reservoir sampling, without entropy decoding.

reservoir-sampling.py

"""
Asymptotically optimal entropy usage for reservoir sampling.
"""

import random

from collections import namedtuple

# Baseline implementation: split odd/even, then pick the last odd element specially, with p = 1/n
def recursive_baseline(items):
    mod_2 = random.randrange(0, 2)
    len = 0
    subselection = []
    last_item = None
    count = 0
    for i, item in enumerate(items):
        count += 1
        last_item = item
        if (i % 2) == mod_2:
            subselection.append(item)
    if count == 0:
        return None
    # If we have an odd number of items, the last item is picked with p = 1 / n
    if (count % 2) == 1:
        # Take it out of the subselected array if it was just put in there.
        if mod_2 == 0:
            subselection.pop()
        if random.randrange(0, count) == count - 1:
            return last_item
    # Otherwise, either the number of items is even, or the last item wasn't chosen.
    # If even:
    #   any item is passed to this recursive call with p_0 = 1/2, and then 
    #   picked with p_1 = 1 / (n / 2) by induction, for p = p_0 * p_1 = 1/n.
    # If odd:
    #  any item in the even-length prefix is passed to this recursive call with
    #  p_0 = (n - 1) / n * 1/2 = (n - 1) / 2 * 1 / n, then picked with
    #  p_1 = 2 / (n - 1), for p = p_0 * p_1 = 1 / n.
    return recursive_baseline(subselection)


"""
In [59]: counts = {0:0, 1:0, 2:0}

In [60]: for _ in range(1200000):
    ...:     counts[recursive_baseline([0, 1, 2])] += 1
    ...:     

In [61]: counts
Out[61]: {0: 399438, 1: 399793, 2: 400769}
"""


# Same thing, but with explicit coroutines (with data direction that's
# somewhat awkward to express with iterators).
Level = namedtuple('Level',  ['mod_2', 'count', 'buffer'])


def add_item_to_reservoir(item, levels, i):
    assert i <= len(levels)
    if i == len(levels):  # new level
        levels.append(Level(random.randrange(0, 2), [0], []))
    level = levels[i]
    level.count[0] += 1
    assert len(level.buffer) < 2
    level.buffer.append(item)
    if len(level.buffer) == 2:
        tournament_winner = level.buffer[level.mod_2]
        level.buffer.clear()
        add_item_to_reservoir(tournament_winner, levels, i + 1)

def finalize_choice(levels):
    assert levels[-1].count[0] == 1
    for level in levels:
        count = level.count[0]
        assert (not level.buffer) == ((count %2) == 0)
        if level.buffer and random.randrange(0, count) == count - 1:
            assert (count % 2) == 1
            return level.buffer[0]
    assert False

def iterative_reservoir(items):
    levels = []
    for item in items:
        add_item_to_reservoir(item, levels, 0)
    return finalize_choice(levels)


"""
In [94]: counts = {0:0, 1:0, 2:0}

In [95]: for _ in range(1200000):
    ...:     counts[iterative_reservoir([0, 1, 2])] += 1
    ...:     

In [96]: counts
Out[96]: {0: 399987, 1: 399570, 2: 400443}
"""

# Now we do the same thing as entropy decoder to recycle unused entropy,
# but more simply, since the ranges over which we have to generate uniforms
# are so nicely related.
def finalize_choice_with_reuse(levels):
    assert levels[-1].count[0] == 1
    random_value = random.randrange(0, levels[0].count[0])

    for level in levels:
        # random_value is uniform in [0, count)
        count = level.count[0]
        if level.buffer and random_value == count - 1:
            return level.buffer[0]
        # If count is odd:
        #   random_value < count - 1, so random_value' = (random_value - 1) / 2 is uniform in (count - 1) / 2,
        #   and the next level's count is (count - 1) / 2, so the loop invariant is maintained.
        # Otherwise, random_value' = random_value / 2 is uniform in count / 2, and the next level's count is count / 2.
        random_value = random_value // 2
    assert False


def reusing_reservoir(items):
    levels = []
    for item in items:
        add_item_to_reservoir(item, levels, 0)
    if False:
        print('n levels: %i' % len(levels))
    return finalize_choice_with_reuse(levels)


"""
In [107]: counts = {i:0 for i in range(13)}

In [108]: for _ in range (13 * 100 * 1000):
     ...:     counts[reusing_reservoir(values)] += 1
     ...: 

In [109]: counts
Out[109]: 
{0: 99567,
 1: 100360,
 2: 99932,
 3: 100518,
 4: 99638,
 5: 100196,
 6: 100385,
 7: 100071,
 8: 100087,
 9: 99997,
 10: 99922,
 11: 99948,
 12: 99379}
"""