O(lgN) weighted random python

gistfile1.py

import random
from bisect import bisect

class WeightedRandom2(object):
    def __init__(self, weights):
        # init interval tree
        self._psum = []
        self._wsum = 0
        for k,v in weights.items():
            assert v > 0, "weight must be a positive integer"
            
            self._wsum += v
            self._psum.append(self._wsum)

    def generate(self):
        rnd = random.random() * self._wsum
        return bisect(self._psum, rnd)
        
# Simple test
from collections import Counter
c = Counter()
wr = WeightedRandom2({1:90,2:5,3:5})
for _ in range(20000):
    c[wr.generate()] += 1
print(c)