[算法]蓄水池抽样算法

Dear 丶 2023-06-25 10:20 8阅读 0赞

问题背景：

> 现有一个单链表，要求随机选择链表中的一个节点并返回节点值，并且保持链表中每个节点被选中的概率相同。

刚看到这个问题，很多人肯定会很不屑：这有何难？先求得链表长度，再生成一个随机数，返回对应位置的节点的值不就完事：

# Definition for singly-linked list.
    # class ListNode(object):
    #     def __init__(self, x):
    #         self.val = x
    #         self.next = None
    
    class Solution(object):
    
        def __init__(self, head):
            """
            @param head The linked list's head.
            Note that the head is guaranteed to be not null, so it contains at least one node.
            :type head: ListNodeQ
            """
            self.arr = []
            while head:
                self.arr.append(head.val)
                head = head.next
            
    
        def getRandom(self):
            """
            Returns a random node's value.
            :rtype: int
            """
            return random.choice(self.arr)
            
    
    # Your Solution object will be instantiated and called as such:
    # obj = Solution(head)
    # param_1 = obj.getRandom()

其实这是LeetCode 第382题，其还有进阶要求：

> 如果链表十分大且长度未知，如何解决这个问题？你能否使用常数级空间复杂度实现？

这么一来这个问题就有点意思了，绞尽脑汁也想不出来解法了。怎么办？这时候就需要用到今天的主角**蓄水池抽样算法**了 。先看看算法伪代码：

# 问题背景：假如现在有海量数据，无法一次将其读入内存，
    # 数据长度也不可得知，暂且用N表示其长度。
    # 现在需要从中抽样k个数据，要求抽样得到的k个数据，每个数据被选中的概率均为k/N.
    
    def sample(Data, k):
        # 先将Data中的前k个数据放入蓄水池pool
        pool = [Data[i] for i in range(k)]
        
        # 从第k+1个数据开始遍历(由于我们是从0开始的，因此索引为k的数就是第k+1个数了)
        for i in range(k, len(Data)):
            # 生成一个随机数，范围为[0,i]
            r = random.randint(0, i)
            
            # 如果生成的随机数落在[0,k)之间，则将当前位置的数据替换掉pool中索引为r+1的数
            if r < k:
                pool[r] = Data[i]
        
        # 返回pool即为所求
        return pool

看完这个伪代码，你是不是会一脸鄙夷：嗯？这么简单的算法就能符合上面的要求？

其实一开始我也是不信的，但是看完了严谨的（让人头大的）证明之后，我只能说：想出这个算法的人真是个人才！下面进行证明：

> 首先对初始就被选中的k个数进行证明（1≤i≤k）：
> 
> 假设第i个数被选中且i≤i小于等于k（即i号球在一开始的时候是在pool中的），那么在选择第k+1个数之前，第i个数被选中的概率为1。在选择第k+1个数的时候，第k+1个数有k/(k+1)的概率被选中，同时在pool中的k个数会有一个被替换（假设为第i个数），那么第i个数被替换掉的概率为(1/k)\*(k/k+1)=1/(k+1)，相反的，第i个数留在pool中的概率为k/(k+1)。
> 
> 同理，在选择第k+2个数的时候，第i个数被替换掉的概率为1/(k+2)，其留在pool中的概率为(k+1)/(k+2)。
> 
> 以此类推，到第N个数的时候，第i个数留在pool中的概率为(N-1)/N。
> 
> 因此第i个数最终留在pool中的概率为：![\\frac\{k\}\{k+1\}\\frac\{k+1\}\{k+2\}...\\frac\{N-1\}\{N\}=\\frac\{k\}\{N\}][frac_k_k_1_frac_k_1_k_2_..._frac_N-1_N_frac_k_N]
> 
> 接着证明i＞k的时候：
> 
> 假设第i个数被选中且k＜i≤N，那么在选择第i个数时，第i个数被选进pool的概率为k/i。在选择第i+1个数的时候，将第i+1个数选进pool并且同时替换第i个数的概率为(1/k)\*(k/i+1)=1/(i+1)，因此第i个数留下来的概率为i/(1+i)，因此从第i个数被选中到第i+1个数的过程中，第i个数仍然留在pool中的概率为：![\\frac\{k\}\{i\}\\frac\{i\}\{i+1\}=\\frac\{k\}\{i+1\}][frac_k_i_frac_i_i_1_frac_k_i_1]
> 
> 同理，在第i+2个数的时候，第i个数仍然留在pool中的概率为：![\\frac\{k\}\{i+2\}][frac_k_i_2]
> 
> 以此类推，到第N个数的时候，第i个数仍然留在pool中的概率为：![\\frac\{k\}\{N\}][frac_k_N]
> 
> 因此第i个数最终留在pool中的概率为![\\frac\{k\}\{N\}][frac_k_N]

最后，使用蓄水池抽样算法解决这个问题的代码如下：

# Definition for singly-linked list.
    # class ListNode(object):
    #     def __init__(self, x):
    #         self.val = x
    #         self.next = None
    
    class Solution(object):
    
        def __init__(self, head):
            """
            @param head The linked list's head.
            Note that the head is guaranteed to be not null, so it contains at least one node.
            :type head: ListNodeQ
            """
            self.head = head
            
    
        def getRandom(self):
            """
            Returns a random node's value.
            :rtype: int
            """
            ans, node, index = self.head, self.head.next, 1
            while node:
                if random.randint(0, index) == 0:
                    ans = node
                node = node.next
                index += 1
            return ans.val
            
    
    
    # Your Solution object will be instantiated and called as such:
    # obj = Solution(head)
    # param_1 = obj.getRandom()

（完）

[frac_k_k_1_frac_k_1_k_2_..._frac_N-1_N_frac_k_N]: https://private.codecogs.com/gif.latex?%5Cfrac%7Bk%7D%7Bk&plus;1%7D%5Cfrac%7Bk&plus;1%7D%7Bk&plus;2%7D...%5Cfrac%7BN-1%7D%7BN%7D%3D%5Cfrac%7Bk%7D%7BN%7D
[frac_k_i_frac_i_i_1_frac_k_i_1]: https://private.codecogs.com/gif.latex?%5Cfrac%7Bk%7D%7Bi%7D%5Cfrac%7Bi%7D%7Bi&plus;1%7D%3D%5Cfrac%7Bk%7D%7Bi&plus;1%7D
[frac_k_i_2]: https://private.codecogs.com/gif.latex?%5Cfrac%7Bk%7D%7Bi&plus;2%7D
[frac_k_N]: https://private.codecogs.com/gif.latex?%5Cfrac%7Bk%7D%7BN%7D