python で重み付きrandom choice

numpy 使うと楽。

import numpy as np

d = [("foo", 98), ("bar", 844), ("baz", 1028)]
a, w = zip(*d)
print(a, w)
# -> ('foo', 'bar', 'baz') (98, 844, 1028)

w2 = np.array(w) / sum(w)
v = np.random.choice(a, p=w2)
print(v)
# -> baz


# 確認
from collections import Counter
c = [ np.random.choice(a, p=w2) for i in range(sum(w)) ]
print(Counter(c))
# -> Counter({'baz': 1051, 'bar': 824, 'foo': 95})

だいたい合ってる。

numpy 使わない版

これだけのために numpy 入れるのもなぁ…な場合は bisect 使うといいかも。

import bisect
import itertools
import random

d = [("foo", 98), ("bar", 844), ("baz", 1028)]
a, w = zip(*d)
print(a, w)

ac = list(itertools.accumulate(w))
f = lambda : bisect.bisect(ac, random.random() * ac[-1])
v = a[f()]
print(v)

# 確認
from collections import Counter
c = [ a[f()] for i in range(sum(w)) ]
print(Counter(c))
# -> Counter({'baz': 1010, 'bar': 851, 'foo': 109})

だいたい合ってる。