正好遇到一个需求,是根据popularity的平滑结果采样,从大约300万个候选集里每次采样100个,重复300万次。
用np.random.choice()性能跟不上。
试了一下这篇文章的算法,稍微提高了一点。
引用自:
https://lotabout.me/2018/Weighted-Random-Sampling/
一个集合里有 n 个元素,每个元素有不同的权重,现在要不放回地随机抽取 m 个元素,每个元素被抽中的概率为元素的权重占总权重的比例。要怎么做呢?
现在考虑只抽取一个元素,假设权重之和为 1 。我们可以从 [0, 1] 中随机得到一个权重,假设为 0.71 ,而后从第一个元素开始,不断累加它们的权重,直到有一个元素的累加权重包含 0.71 ,则选取该元素。下面是个示意图:
加权随机采样 (Weighted Random Sampling)
要选取 m 个元素,则可以按上面的方法先选取一个,将该元素从集合中去除,再反复按上面的方法抽取剩余的元素。这种方法的复杂度是 O(mn) ,并且将元素从集合中删除其实不太方便实现。
当然,最重要的是这个算法需要多次遍历数据,不适合用在流处理的场景中。
Algorithm A 是论文 Weighted Random Sampling 中提出的,步骤如下:
- 对于集合 V V V 中的元素 v i ∈ V v_i \in V vi∈V,选取均匀分布的随机数 u i = r a n d ( 0 , 1 ) u_i = rand(0, 1) ui=rand(0,1) ,计算元素的特征 k i = u i ( 1 / w i ) k_i = u_i^{(1/w_i)} ki=ui(1/wi)
- 将集合按 k i k_i ki 排序,选取前 m m m 大的元素。
算法的正确性在作者 2006 年的论文 Weighted random sampling with a reservoir 里给了详细的证明。论文中给出了算法的两个变种 A-Res 与 A-ExpJ,它们都能在一次扫描中得到 m 个样本。非常适合在流处理的场合中。
A-Res(Algorithm A With a Reservoir) 是 Algorithm 的“蓄水池”版本,即维护含有 m 个元素的结果集,对每个新元素尝试去替换结果集中权重最小的元素。步骤如下:
- 将集合 V V V 的前 m m m 个元素放入结果集合 R R R。
- 对于结果集里的每个元素,计算特征值 k i = u i ( 1 / w i ) k_i = u_i^{(1/w_i)} ki=ui(1/wi),其中 u i = r a n d ( 0 , 1 ) u_i = rand(0, 1) ui=rand(0,1)
- 对 i = m + 1 , m + 2 , … , n i = m+1, m+2, \dots, n i=m+1,m+2,…,n 重复步骤 4 ~ 6
- 将结果集中最小的特征 k k k 作为当前的阈值 T T T
- 对于元素 v i v_i vi,计算特征 k i = u i ( 1 / w i ) k_i = u_i^{(1/w_i)} ki=ui(1/wi),其中 u i = r a n d ( 0 , 1 ) u_i = rand(0, 1) ui=rand(0,1)
- 如果 k i > T k_i > T ki>T 则将 R R R 中拥有最小 k k k 值的元素替换成 v i v_i vi。
论文证明了如果权重 w i w_i wi 是一般连续分布上的随机变量,则上面的算法中插入 R R R 的次数为 O ( m log ( n m ) ) O(m \log(\frac{n}{m})) O(mlog(mn))。该算法用 Python 实现如下:
import heapq
import random
def a_res(samples, m):
"""
:samples: [(item, weight), ...]
:k: number of selected items
:returns: [(item, weight), ...]
"""
heap = [] # [(new_weight, item), ...]
for sample in samples:
wi = sample[1]
ui = random.uniform(0, 1)
ki = ui ** (1/wi)
if len(heap) < m:
heapq.heappush(heap, (ki, sample))
elif ki > heap[0][0]:
heapq.heappush(heap, (ki, sample))
if len(heap) > m:
heapq.heappop(heap)
return [item[1] for item in heap]
A-Res 需要对每个元素产生一个随机数,而生成高质量的随机数有可能会有较大的性能开销,,所以论文中给出了 A-ExpJ 算法,能将随机数的生成量从 O ( n ) O(n) O(n) 减少到 O ( m log ( n m ) ) ) O(m\log(\frac{n}{m}))) O(mlog(mn)))。
从步骤上看,A-ExpJ 算法很像我们最开始提出的简单版本,设定一个阈值并跳过一些元素。具体步骤如下:
- 将集合 V V V 的前 m m m 个元素放入结果集合 R R R。
- 对于结果集里的每个元素,计算特征值 k i = u i ( 1 / w i ) k_i = u_i^{(1/w_i)} ki=ui(1/wi),其中 u i = r a n d ( 0 , 1 ) u_i = rand(0, 1) ui=rand(0,1)
- 将 R R R 中小最的特征值记为阈值 T w T_w Tw
- 对剩下的元素重复步骤 5 ~ 9
- 令 r = r a n d ( 0 , 1 ) r = rand(0, 1) r=rand(0,1) 且 X w = log ( r ) / log ( T w ) X_w = \log( r )/\log(T_w) Xw=log(r)/log(Tw)
- 从当前元素
v
c
v_c
vc 开始跳过元素,直到遇到元素
v
i
v_i
vi,满足
w c + w c + 1 + ⋯ + w i − 1 < X w ≤ w c + w c + 1 + ⋯ + w i − 1 + w i w_c + w_{c+1} + \dots + w_{i-1} \lt X_w \le w_c + w_{c+1} + \dots + w_{i-1} + w_{i} wc+wc+1+⋯+wi−1<Xw≤wc+wc+1+⋯+wi−1+wi - 使用 v i v_i vi 替换 R R R 中特征值最小的元素。
- 令 t w = T w w i t_w = T_w^{w_i} tw=Twwi, r 2 = r a n d ( t w , 1 ) r2 = rand(t_w, 1) r2=rand(tw,1), v i v_i vi 的特征 k i = r 2 ( 1 / w i ) k_i = r_2^{(1/w_i)} ki=r2(1/wi)
- 令新的阈值 T w T_w Tw 为此时 R R R 中的最小特征值。
Python 实现如下:
import heapq
import numpy as np
def a_expj(samples, m):
"""
:samples: [(item, weight), ...]
:k: number of selected items
:returns: [(item, weight), ...]
"""
heap = [] # [(new_weight, item), ...]
Xw = None
Tw = 0
w_acc = 0
for sample in samples:
if len(heap) < m:
wi = sample[1]
ui = np.random.uniform(0, 1)
ki = ui ** (1/wi)
heapq.heappush(heap, (ki, sample))
continue
if w_acc == 0:
Tw = heap[0][0]
r = np.random.uniform(0, 1)
Xw = np.log(r)/np.log(Tw)
wi = sample[1]
if w_acc + wi < Xw:
w_acc += wi
continue
else:
w_acc = 0
tw = Tw ** wi
r2 = np.random.uniform(tw, 1)
ki = r2 ** (1/wi)
heapq.heappop(heap)
heapq.heappush(heap, (ki, sample))
#print(heap)
return [item[1] for item in heap]
我们用多次采样的方式来尝试验证算法的正确性。下面代码中为 a 、 b 、 c 等元素赋予了不同的权重,采样 10 万次后计算被采样的次数与元素 a 被采样次数的比值。
overall = [('a', 10), ('b', 20), ('c', 50), ('d', 100), ('e', 200)]
def test_weighted_sampling(func, k):
stat = {}
for i in range(100000):
sampled = func(overall, k)
for item in sampled:
if item[0] not in stat:
stat[item[0]] = 0
stat[item[0]] += 1
total = stat['a']
for a in stat:
stat[a] = float(stat[a])/float(total)
print(stat)
首先验证 A-Res 算法:
#说明
#需要的输入是形状是
#samples = [(id,weight),(id,weight)...]
#可以用zip达到我们的目的
samples = list(zip(candidates,probs))
#测试开始
test_weighted_sampling(a_res, 1)
test_weighted_sampling(a_res, 2)
test_weighted_sampling(a_res, 3)
test_weighted_sampling(a_res, 4)
test_weighted_sampling(a_res, 5)
output
{'e': 19.54951600893522, 'd': 9.864110201042442, 'c': 4.842889054355919, 'a': 1.0, 'b': 1.973566641846612}
{'b': 2.0223285486443383, 'e': 12.17949833260838, 'd': 8.95287806292591, 'c': 4.843410178338408, 'a': 1.0}
{'a': 1.0, 'e': 6.166443722530097, 'd': 5.597171794381808, 'b': 1.9579591056755208, 'c': 4.387922797630423}
{'b': 1.8358898492044953, 'e': 2.5878688779880092, 'c': 2.4081341327311896, 'd': 2.549897479820395, 'a': 1.0}
{'a': 1.0, 'd': 1.0, 'c': 1.0, 'b': 1.0, 'e': 1.0}
看到,在采样一个元素时, b 被采样到的次数约为 a 的 2 倍,而 e 则约为 20 倍,与 overall 数组中指定的权重一致。而采样 5 个元素时,所有元素都会被选中。
同理验证 A-ExpJ 算法:
test_weighted_sampling(a_expj, 1)
test_weighted_sampling(a_expj, 2)
test_weighted_sampling(a_expj, 3)
test_weighted_sampling(a_expj, 4)
test_weighted_sampling(a_expj, 5)
output
{'e': 19.78311444652908, 'c': 4.915572232645403, 'd': 9.840900562851782, 'a': 1.0, 'b': 1.9838649155722325}
{'e': 11.831543244771057, 'c': 4.709157716223856, 'b': 1.9720180893159978, 'd': 8.75183719615602, 'a': 1.0}
{'d': 5.496249062265567, 'c': 4.280007501875469, 'e': 6.046324081020255, 'b': 1.9321080270067517, 'a': 1.0}
{'a': 1.0, 'd': 2.5883654175335105, 'c': 2.440760540383957, 'e': 2.62591841571643, 'b': 1.8787559581808126}
{'a': 1.0, 'd': 1.0, 'c': 1.0, 'b': 1.0, 'e': 1.0}
与 A-Res 的结果类似。
文章中介绍了 A-Res 与 A-ExpJ 两种算法,按照步骤用 Python 实现了一个简单的版本,最后用采样的方式验证了算法的正确性。
加权随机采样本身不难,但如果需要在一次扫描中完成就不容易了。难以想像上面的算法直到 2006 年才提出。算法本身如此之简单,也让不感叹数学与概率的精妙。