依概率不放回随机抽样算法

起因

前段时间帮同学写了下ER网络和BA网络。
其中BA网络要根据度占整个图的度的比例作为概率。
于是就写了个按概率随机抽数的函数pick。

思路

当时想法很简单。就是把数轴分成几块，再随机抽点。

比如数组[1,2,3]，就生成区间[1,6]的随机整数。
若随机数为5就认为选中第三个数。

问题

上面方法抽一个数很有效，但如果是抽n个数性能就下降的很厉害。越抽到后面重复的概率越大。

我刚开始用洗牌算法，抽到第i个数，就和位置i的数字交换，再对后面数字依概率取数。这样虽然解决了重复取数的问题，但是效率仍然不高。因为每次都要重新计算前缀和。

后来又去改进蓄水池算法。但似乎想不到非等概率的算法。

最终算法

最后还是采用把概率设为0的方式。既避免重复取数，也方便更新前缀和。
前缀和用树状数组维护，查询修改都是$log(n)$
判断是哪个区间用二分查找，二分查找本来复杂度为$log(n)$，里面又有个查询前缀和。总复杂度$(log(n){)}^2$

下面举个例子好理解点
比如概率比例为[1,2,3],前缀和为[1,3,6]
如果第一次选到第二个数。那么就变成 概率比例为[1,0,3],前缀和为[1,1,4]，二分法对前缀和每次查找第一个大于等于随机数的位置就既能保证不重复又能保证时间复杂度了。

代码

from scipy.stats import ks_2samp
import numpy as np
import random
from typing import List


class TreeArray:
    def __init__(self, arr: List[int]) -> None:
        self.len = len(arr)+1
        self.arr = [0]*self.len
        for i in range(0, self.len-1):
            self.add(i, arr[i])

    def add(self, i: int, d: int) -> None:
        '''i位置加上d'''
        i += 1
        while i < self.len:
            self.arr[i] += d
            i += (i & (-i))

    def sum(self, i: int) -> int:
        '''获取到前i个数的和'''
        res = 0
        while i > 0:
            res += self.arr[i]
            i -= (i & (-i))
        return res
    pass

def bsearch(arr: TreeArray, k: int) -> int:
    '''返回大于等于k的下标,若多个符合返回最小的'''
    l, r = -1, arr.len-1
    while l+1 < r:
        m = int((l+r)/2)
        if arr.sum(m+1) < k:
            l = m
        else:
            r = m
    return r


def pick(ps: List[int], n: int) -> List[int]:
    '''以特定概率比例ps随机选取n个数'''
    section,  res = TreeArray(ps), [None]*n
    cur_sum = section.sum(len(ps))
    for i in range(n):
        # [1,cur_sum]
        x = random.randint(1, cur_sum)
        j = bsearch(section, x)
        res[i], cur_sum = j, cur_sum-ps[j]
        section.add(j, -ps[j])
    return res

测试

from scipy.stats import ks_2samp
import numpy as np

n = 500
s = 5
p = np.array([5, 3, 2, 5, 2, 1, 8, 3, 2, 1])
p1 = p/sum(p)
ids = []
for i in range(len(p)):
    ids.append(i)

cnt = [0]*len(p)
cnt1 = [0]*len(p)
for i in range(n):
    ans = pick(p, s)
    for i in range(s):
        cnt[ans[i]] += 1
    ans = np.random.choice(ids, size=s, replace=False, p=p1)
    for i in range(s):
        cnt1[ans[i]] += 1
print(ks_2samp(cnt1, cnt))

这里我抽了500轮。KS-检验判断是否符合概率分布。
这里的对拍函数是numpy库里的np.random.choice。

可以看到pvalue很高啊。
根据源码的注释可以看出pvalue愈高或statistic越低,两个分布的相似度越高。

If the KS statistic is small or the p-value is high, then we cannot
reject the hypothesis that the distributions of the two samples are
the same.

插曲

你以为到这完了？其实并没有。
我增大数据量时，奇怪的事发生了。
按理来说数据量越大，两个分布就越相似。但我测出来的结果刚好相反。
它的趋势不是单调的，
n取500时
KstestResult(statistic=0.2, pvalue=0.9944575548290717)
n取5000时
KstestResult(statistic=0.1, pvalue=1.0)
n取50000时
KstestResult(statistic=0.3, pvalue=0.7869297884777761)

可以发现取5000时效果最好。
后来我又去查了写资料发现ks不适合离散型分布检验
具体原因参阅这里