load_balancer.py

import random
from collections import defaultdict

class LoadBalancer:
    def __init__(self, targets, probabilities):
        self._targets = targets.copy()
        n = len(targets)
        self._n = n
        p = probabilities.copy()
        self._alias = [0] * n
        self._prob = [0] * n
        small = []
        large = []
        for i in range(n):
            p[i] = p[i] * n
            if p[i] < 1:
                small.append(i)
            else:
                large.append(i)
        while len(small) > 0 and len(large) > 0:
            l = small.pop()
            g = large.pop()
            self._prob[l] = p[l]
            self._alias[l] = g
            p[g] = p[g] + p[l] - 1
            if p[g] < 1:
                small.append(g)
            else:
                large.append(g)
        for i in small:
            self._prob[i] = 1
        for i in large:
            self._prob[i] = 1

    def pick(self) -> str:
        slot = random.randint(0, self._n-1)
        toss = random.random()
        if toss < self._prob[slot]:
            return self._targets[slot]
        else:
            return self._targets[self._alias[slot]]


def main():
    targets = ["foo", "bar", "baz"]
    prob = [0.1, 0.3, 0.6]
    lb = LoadBalancer(targets, prob)
    results = defaultdict(int)
    for i in range(1000):
        pick = lb.pick()
        results[pick] += 1
    for k in targets:
        print("%s: %d" % (k, results[k]))


if __name__ == "__main__":
    main()