#!/usr/bin/env python
# This is a version of the Alias Method for weighted random selection
# that uses big and small indices instead of generators. It's not
# really faster, but it is more suitable to porting to languages like
# C that don't have generators.
#
# Copyright 2019 Bruce Hill.
# This code is released under the MIT license.
# For details, see: https://opensource.org/licenses/MIT

def prepare_aliased_randomizer(weights):
    N = len(weights)
    avg = sum(weights)/N
    aliases = [(1, None)]*N
    small_i = 0
    while small_i < N and weights[small_i] >= avg:
        small_i += 1
    if small_i < N: # If all weights are the same, nothing to do
        small = (small_i, weights[small_i]/avg)
        big_i = 0
        while big_i < N and weights[big_i] < avg:
            big_i += 1
        assert(big_i < N) # At least one weight is >= avg
        big = (big_i, weights[big_i]/avg)
        while big and small:
            aliases[small[0]] = (small[1], big[0])
            big = (big[0], big[1] - (1-small[1]))
            if big[1] < 1:
                small = big
                big_i += 1
                while big_i < N and weights[big_i] < avg:
                    big_i += 1
                if big_i >= N: break
                big = (big_i, weights[big_i]/avg)
            else:
                small_i += 1
                while small_i < N and weights[small_i] >= avg:
                    small_i += 1
                if small_i >= N: break
                small = (small_i, weights[small_i]/avg)
    
    def weighted_random():
        r = random()*N
        i = int(r)
        odds, alias = aliases[i]
        return alias if (r-i) > odds else i

    return weighted_random