import numpy as np

def alias_setup(probs):
    K = len(probs)
    q = np.zeros(K)
    J = np.zeros(K, dtype=np.int)

    smaller = []
    larger = []
    for kk, prob in enumerate(probs):
        q[kk] = K * prob         # 保存面积块
        if q[kk] < 1.0:
            smaller.append(kk)   # 面积小于1的元素下标
        else:
            larger.append(kk)    # 面积大于1的元素下标

    while len(smaller) > 0 and len(larger) > 0:
        small = smaller.pop()
        large = larger.pop()
        # J 记录的是填充这一列的下标
        J[small] = large  
        # q 记录的是原来事件的占比
        q[large] = q[large] - (1 - q[small])
        if q[large] < 1.0:
            smaller.append(large)
        else:
            larger.append(large)
    
    return J, q

def alias_draw(J, q):
    K = len(J)
    kk = int(np.floor(np.random.rand() * K))  # 随机选一列
    if np.random.rand() < q[kk]:
        return kk
    else:
        return J[kk]
    
if __name__  == "__main__":
    J, q = alias_setup([1/2, 1/3, 1/12, 1/12])
    print(J, q)   # [0 0 0 1] [1.         0.66666667 0.33333333 0.33333333]
    ret = alias_draw(J, q)