確率を指定してリストからランダムに取得する

仕事で、指定した確率(重み付け)に従ってリストの中からランダムで一つ取得するという仕組みが必要になったのでこんなのを書いてみました。
言語はPythonです。NYSLでどーぞ。

import random

def choose(candidates, probabilities):
    probabilities = [sum(probabilities[:x+1]) for x in range(len(probabilities))]
    if probabilities[-1] > 1.0:
        #確率の合計が100%を超えていた場合は100％になるように調整する
        probabilities = [x/probabilities[-1] for x in probabilities]
    rand = random.random()
    for candidate, probability in zip(candidates, probabilities):
        if rand < probability:
            return candidate
    #どれにも当てはまらなかった場合はNoneを返す
    return None

candidatesに選択したいものを並べたリストを、probabilitiesに確率(重み付け)のリストを渡すと、その確率に沿ってcandidatesからひとつ取得します。
↓choose(['A', 'B', 'O', 'AB'], [0.4, 0.2, 0.3, 0.1])で10万回施行
A : 40037
B : 20089
AB : 9913
O : 29961

probabilitiesの合計が100％未満の場合はNoneが返ることがあります。
↓choose(['foo', 'bar', 'baz'], [0.1, 0.2, 0.3])で10万回施行
foo : 10274
bar : 19905
baz : 29746
None : 40075

逆に100％を超えている場合は、全部で100％になるように確率を調整します。
↓choose(['hoge', 'fuga', 'piyo'], [0.4, 0.6, 1.0])で10万回施行
hoge : 19794
fuga : 30023
piyo : 50183

Tweet