重み付きランダム抽出:Walker's Alias MethodのPython版
python
import numpy as np
def rand_from_prob(wgt, nn=1):
"""Walker's Alias Method"""
if not isinstance(wgt, np.ndarray):
wgt = np.array(wgt)
wsm = sum(wgt)
n = len(wgt)
p = (wgt*n/wsm).tolist()
a, hl = [0] * n, [0] * n
l, h = 0, n-1
for i in range(n):
if p[i] < 1:
hl[l] = i
l += 1
else:
hl[h] = i
h -= 1
while l > 0 and h < n-1:
j, k = hl[l-1], hl[h+1]
a[j] = k
p[k] += p[j] - 1
if p[k] < 1:
hl[l-1] = k
h += 1
else:
l -= 1
rr = np.random.rand(nn) * n
ii = np.int32(np.floor(rr))
rr -= ii
return np.array([i if r < p[i] else a[i] for i, r in zip(ii, rr)])
O(1)なので高速。計算時間の比較。
python
w = np.random.rand(2000)
w /= sum(w)
%time x = rand_from_prob(w, 100000)
>>>
Wall time: 52 ms
%time y = np.random.multinomial(1, w, 100000)
>>>
Wall time: 6.75 s