重み付け乱択は、2段階で改良できます。
まず、1段階目として、リストを1回走査するだけで(重みの合計値を求めながら)
乱択することができます。
import random
def random_weight_choice(L):
choice = None
total = 0
for item, p in L:
total += p
if random.random() * total < p:
choice = item
return choice
def test_random_weight_choice():
from collections import defaultdict
X = [('A', 3), ('B', 2), ('C', 5)]
count = defaultdict(int)
for _ in xrange(100000):
item = random_weight_choice(X)
count[item] += 1
print count
if __name__ == '__main__':
test_random_weight_choice()
この段階では、走査が1パスになっただけで、計算量は O(n) のままです。
同じ重み付けリストから乱択をなんども繰り返す場合は、事前に専用のデータ構造を
作成することで、 O(log n) や O(1) に計算量を落とすことができます。
Python の場合、 bisect というモジュールで2分探索が簡単にできるので、
それを使えば O(log n) の乱択を簡単に実装できます。
http://codepad.org/TTjAK1mV
import random
import bisect
class Choicer(object):
def __init__(self, L):
total = 0
self.probs = []
self.items = []
for item, p in L:
total += p
self.probs.append(total)
self.items.append(item)
def choice(self):
r = random.random() * self.probs[-1]
idx = bisect.bisect(self.probs, r)
return self.items[idx]
def test_random_weight_choice():
from collections import defaultdict
X = [('A', 3), ('B', 2), ('C', 5)]
choicer = Choicer(X)
count = defaultdict(int)
for _ in xrange(100000):
item = choicer.choice()
count[item] += 1
print count
if __name__ == '__main__':
test_random_weight_choice()