聊聊最近面试中遇到的算法题:公平的随机

前几天字节跳动的朋友发现一个很有趣的问题，在这里和读者们分享一下。

题目描述贼简单：设计一个随机抽样的系统，这个系统需要实时的从数据流中读取数据，并提供一个查询接口，这个接口需要从已读取的数据中等概率的选取K个数并返回。

一开始我以为是个设计随机算法的题目，所以给出了一个保存所有已读取数据的思路，然后刚说完“保存全部数据”就被面试官无情打断 —— 数据流很大很大，大到无法保存所有数据。

这样的话，空间复杂度应该不能超过O(K)，只能是从已保存的K份数据和新读入的数据上做文章了。然后和面试官沟通了一些细节：

• O(K)的空间复杂度是否可以? (可以)
• 数据流中的数据总量会超过 int64 吗？(不会超过)

Tips：前面被打断就是因为在思考解决方案前没有和面试官沟通好数据规模。所以应该及时调整心态，多和面试官沟通一下细节。毕竟面试的时候时间和精力都很宝贵，尽量不要犯同样的错误啊老铁。

得到面试官的答案后继续思考。假设刚刚读取到第 N 个数据：

• 如果 N <= K，那么该数据被保留下的概率肯定为 1
• 如果 N > K，那么该数据被保留下的概率应为 K/N

那该如何实现呢？当 N<=K的时候，显而易见直接保留下来就可以了。
接着考虑 N = K+1 时，该元素被保留下的概率为 K/(K+1)。该元素被保留就意味着前K个元素中肯定有一个被淘汰。
那被淘汰的概率是多少呢？有K/(K+1)的概率触发淘汰操作，如果我们随机地选择一个元素，那么每个元素都有1/K的概率被选中。
综上，被淘汰的概率为 K/(K+1)*1/K，即 1/(K+1)。
因为一个元素要么保留下来，要么被淘汰，所以其保留下的概率亦为 K/(K+1)。

下面数学归纳法证明上述操作使得当输入N个元素时每个元素被保留的概率均为K/N (觉得不妥的老铁可以点击“阅读原文”一起讨论)
当 N = K 时，概率为 K/N = 1，显然成立。
当 N = K+1时，概率 为 K/N = K/(K+1)，如上所述也成立。
那么当 N’ = N+1时，第N’个元素以K/N’的概率决定是否保留，此时前N元素中每个元素被保留下来的概率均为 1 - K/N * K/N’ * 1/K = K/N’。
综上可得，当已读取N个元素时，这种操作可使每个元素被保留下的概率均为 K/N。

总结一下，当读取到第N个元素时，先随机一下是否应该保留，被保留的概率为K/N。如果随机到保留，再随机一下要淘汰掉哪个元素即可。

template <typename DataType>
class RandomK {
  uint64_t count;
  uint32_t k;
  vector<DataType> data_list;
 public:
  RandomK(uint32_t k_) : k(k_) {}
  void input(const DataType &data) {
    count += 1;
    // 直接保留
    if (count <= k) {
      data_list.push_back(data);
      return ;
    }
    // 随机产生一个 [0, count-1] 内的数字 x
    // 如果 x ∈ [0, k-1]，那么该数据应该被保留
    if(random(count) < k) {
      // 每个元素都有 1/K 的概率被淘汰
      int pos = random(k);
      data_list[pos] = data;
    }
  }
  const vector<DataType>& getDataList() const {
    return data_list;
  }
};

来源：https://mp.weixin.qq.com/s/dxZc9opSInsirjFK5L2qMg