low_variance_sampling.py

class DiscreteDistribution:
    def __init__(self):
        self.total_weight = 0
        self.points = {}

    def add(self, point, weight):
        assert point not in self.points
        self.points[point] = weight
        self.total_weight += weight

    def remove(self, point):
        self.total_weight -= self.points[point]
        del self.points[point]

    def sample(self, num_samples):
        items = iter(self.points.items())
        point, cum_weight = next(items)
        x = 0.0
        dx = self.total_weight / num_samples
        for i in range(num_samples):
            yield point
            x += dx
            while x >= cum_weight:
                point, weight = next(items)
                cum_weight += weight

dist = DiscreteDistribution()
dist.add('a', 10)
dist.add('b', 5)
dist.add('c', 3)
print(list(dist.sample(6))) # ['a', 'a', 'a', 'a', 'b', 'c']
dist.remove('a')
print(list(dist.sample(6))) # ['b', 'b', 'b', 'b', 'c', 'c']