数据结构与算法：趣味算法 —— 蓄水池抽样

# Python常用内置随机函数
import random

# random产生[0,1)之间的浮点数
random.random()
0.9866624765042807

# 产生[a,b]之间的整数
random.randint(1,10)
4

# 产生range(start,end,step)中的随机整数
random.randrange(2,10,2)
6

# 从指定序列中随机选取一个元素
random.choice([1,2,5])
5

蓄水池抽样算法

原理

问题：给定一个长度为n的数组，从中随机抽取k个元素。要求每个元素被选取的概率相等，且时间复杂度为O(n)。
思路：维护一个长度为k的蓄水池，首先将前k个元素放入蓄水池，然后从第k+1个元素开始向后遍历，遍历到第i个元素时，以k/i的概率选中第i个元素，随机替换掉蓄水池中任一元素，直至遍历完成，蓄水池中的k个元素就是最终抽样结果。
证明：蓄水池抽样算法使得每个元素被抽到的概率均为k/n。我们只需要证明遍历完第i个元素时，前i个元素中每个元素被抽取到概率均为k/i：
- 当i=k时，前i个元素每个被抽取到的概率为$1=\frac{k}{i}$
- 假设遍历完第i个元素后，前i个元素每个元素被抽取到的概率为$\frac{k}{i}$，现在考虑第i+1个元素，以$\frac{k}{i+1}$的概率选中第i+1个元素，替换蓄水池中随机一个元素x：
  - 对于第i+1个元素来说，放入蓄水池的概率为$\frac{k}{i+1}$
  - 对于前i个元素来说，遍历第i+1个元素前，每个元素在蓄水池中的概率为$\frac{k}{i}$，遍历完第i+1个元素后仍在蓄水池中的概率为：$\frac{k}{i}\cdot (1-\frac{k}{i+1}\cdot \frac{1}{k})=\frac{k}{i+1}$
- 当i=n时，前n个元素被抽取到的概率均为$\frac{k}{n}$
实现：

import random
def reservoid_sample(a,k):
    res = []
    for i,num in enumerate(a):
        if i < k:
            res.append(num)
        else:
            if random.random() < 1. * k / (i+1):
                res[random.randint(0,k-1)] = num
    return res if i >= k - 1 else False

例题

LeetCode[398] 随机抽样

问题：给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。您可以假设给定的数字一定存在于数组中
思路：只对等于给定数字的元素计数
代码：

import random
class Solution(object):
    def __init__(self, nums):
        """
        :type nums: List[int]
        :type numsSize: int
        """
        self.nums = nums
        

    def pick(self, target):
        """
        :type target: int
        :rtype: int
        """
        count = 0
        for i,num in enumerate(self.nums):
            if num == target: 
                if count == 0:
                    pre = i
                    count += 1
                else:
                    count += 1
                    if random.random() < 1./count:
                        pre = i
        return pre

[LeetCode 528] 按权重随机抽样

问题：

给定一个正整数数组 w ，其中 w[i] 代表位置 i 的权重，请写一个函数 pickIndex ，它可以随机地获取位置 i，选取位置 i 的概率与 w[i] 成正比。

说明:

1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex 将被调用不超过 10000 次

思路：

思路：
求数组各位置之前的累加和
在总和范围内产生随机数
通过二分查找确定随机数所属区间(第一个大于随机数的位置索引)

代码：

import random
import bisect
class Solution(object):
    def __init__(self, w):
        """
        :type w: List[int]
        """
        self.w = w
        for i in xrange(1, len(w)):
            self.w[i] += self.w[i-1]

    def pickIndex(self):
        """
        :rtype: int
        """
        val = random.randint(1, self.w[-1])
        return bisect.bisect_left(self.w, val)