Problem #15 [Medium]

This problem was asked by Facebook.

Given a stream of elements too large to store in memory, pick a random element from the stream with uniform probability.
这是个经典面试问题：蓄水池抽样。如果没有内存限制，我们可以把stream存到一个列表中，然后从[0, size-1]中随机选择元素，空间复杂度为O(N)。

import random
random_element = random.choice(my_list)

print(random_element)

相反，让我们尝试使用循环不变量来解决。在循环的第 i 次迭代中选择随机元素时，假设我们已经从 [0, i - 1] 中统一选择了一个元素。为了保持循环不变，我们需要以 1 / (i + 1) 的机会选择第 i 个元素作为新的随机元素。对于 i = 0 的基本情况，假设随机元素是第一个。然后我们知道它有效，因为

对于 i = 0, we would’ve picked uniformly from [0, 0].
对于 i > 0, before the loop began, any element K in [0, i - 1] had 1 / i chance of being chosen as the random element. We want K to have 1 / (i + 1) chance of being chosen after the iteration. This is the case since the chance of having being chosen already but not getting swapped with the ith element is 1 / i * (1 - (1 / (i + 1))) which is 1 / i * i / (i + 1) or 1 / (i + 1)
代码如下：

import random

def pick(big_stream):
    random_element = None

    for i, e in enumerate(big_stream):
        if i == 0:
            random_element = e
        if random.randint(1, i + 1) == 1:
            random_element = e
    return random_element

空间复杂度为O(1)。