[GO] Weighted Random Extraction: Walker's Alias Method (Python)

Python version of Weighted Random Extraction: Walker's Alias Method

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)])

High speed because it is O (1). Comparison of calculation time.

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