Alias Method

alias.go

package randext

import (
	"math/rand"
)

// Alias Method based on:
// http://www.keithschwarz.com/darts-dice-coins/
// http://www.keithschwarz.com/interesting/code/?dir=alias-method
type Alias struct {
	probabilities []float64
	alias         []int
}

// Constructs a new Alias type to sample from a discrete distribution and
// hand back outcomes based on the probability distribution.
//
// Given as input a list of probabilities corresponding to outcomes 0, 1,
// ..., n - 1, along with the random number generator that should be used
// as the underlying generator, this constructor creates the probability
// and alias tables needed to efficiently sample from this distribution.
func NewAlias(probabilities ...float64) *Alias {
	{
		tmp := make([]float64, len(probabilities))
		for i, p := range probabilities {
			if p < 0.0 {
				p = 0.0
			} else if p > 1.0 {
				p = 1.0
			}
			tmp[i] = p
		}
		probabilities = tmp
	}
	alias := &Alias{make([]float64, len(probabilities)), make([]int, len(probabilities))}

	average := 1.0 / float64(len(probabilities))
	small, large := make([]int, 0), make([]int, 0)

	for i, p := range probabilities {
		if p >= average {
			large = append(large, i)
		} else {
			small = append(small, i)
		}
	}

	for len(small) > 0 && len(large) > 0 {
		less := small[len(small)-1]
		small = small[:len(small)-1]
		more := large[len(large)-1]
		large = large[:len(large)-1]

		alias.probabilities[less] = probabilities[less] * float64(len(probabilities))
		alias.alias[less] = more

		probabilities[more] = (probabilities[more] + probabilities[less]) - average

		if probabilities[more] >= average {
			large = append(large, more)
		} else {
			small = append(small, more)
		}
	}

	for len(small) > 0 {
		less := small[len(small)-1]
		small = small[:len(small)-1]
		alias.probabilities[less] = 1.0
	}
	for len(large) > 0 {
		more := large[len(large)-1]
		large = large[:len(large)-1]
		alias.probabilities[more] = 1.0
	}

	return alias
}

// Samples a value from the underlying distribution
func (alias *Alias) Next(r *rand.Rand) int {
	column := r.Intn(len(alias.probabilities))
	if r.Float64() < alias.probabilities[column] {
		return column
	} else {
		return alias.alias[column]
	}
}

alias_test.go

package randext

import (
	"math"
	"math/rand"
	"testing"
)

func TestAlias(t *testing.T) {

	p := []float64{0.0, 0.1, 0.8, 0.1}
	alias := NewAlias(p...)

	sum := [4]int{}
	r := rand.New(rand.NewSource(777))
	count := 100000
	for i := 0; i < count; i++ {
		sum[alias.Next(r)]++
	}
	for i := 0; i < len(sum); i++ {
		actual := float64(sum[i]) / float64(count)
		expected := p[i]
		if math.Abs(actual-expected) > 0.01 {
			t.Fail()
		}
	}
}

AliasMethod.java

package net.monoid.util;

import java.util.*;

/**
 * based on: 
 * http://www.keithschwarz.com/darts-dice-coins/ 
 * http://www.keithschwarz.com/interesting/code/?dir=alias-method 
 */
public final class AliasMethod {
    private final Random random;
    private final int[] alias;
    private final double[] probability;

    /**
     * Constructs a new AliasMethod to sample from a discrete distribution and
     * hand back outcomes based on the probability distribution.
     * <p>
     * Given as input a list of probabilities corresponding to outcomes 0, 1,
     * ..., n - 1, this constructor creates the probability and alias tables
     * needed to efficiently sample from this distribution.
     *
     * @param probabilities The list of probabilities.
     */
    public AliasMethod(double... probabilities) {
        this(probabilities.clone(), new Random());
    }

    /**
     * Constructs a new AliasMethod to sample from a discrete distribution and
     * hand back outcomes based on the probability distribution.
     * <p>
     * Given as input a list of probabilities corresponding to outcomes 0, 1,
     * ..., n - 1, along with the random number generator that should be used
     * as the underlying generator, this constructor creates the probability 
     * and alias tables needed to efficiently sample from this distribution.
     *
     * @param probabilities The list of probabilities.
     * @param seed The random number generator
     */
    public AliasMethod(double[] probabilities, long seed) {
        this(probabilities.clone(), new Random(seed));
    }
    
    private AliasMethod(double[] probabilities, Random random) {
        if (probabilities.length == 0) {
            throw new IllegalArgumentException("Probability vector must be nonempty.");
        }

        this.random = random;
        
        probability = new double[probabilities.length];
        alias = new int[probabilities.length];

        double average = 1.0 / probabilities.length;

        Deque<Integer> small = new ArrayDeque<Integer>();
        Deque<Integer> large = new ArrayDeque<Integer>();

        for (int i = 0; i < probabilities.length; ++i) {
            if (probabilities[i] >= average) {
                large.add(i);
            } else {
                small.add(i);
            }
        }

        while (!small.isEmpty() && !large.isEmpty()) {
            int less = small.removeLast();
            int more = large.removeLast();

            probability[less] = probabilities[less] * probabilities.length;
            alias[less] = more;

            probabilities[more] = (probabilities[more] + probabilities[less]) - average;

            if (probabilities[more] >= average) {
                large.add(more);
            } else {
                small.add(more);
            }
        }

        while (!small.isEmpty()) {
            probability[small.removeLast()] = 1.0;
        }
        while (!large.isEmpty()) {
            probability[large.removeLast()] = 1.0;
        }
    }

    /**
     * Samples a value from the underlying distribution.
     *
     * @return A random value sampled from the underlying distribution.
     */
    public int next() {
        int column = random.nextInt(probability.length);
        boolean coinToss = random.nextDouble() < probability[column];
        return coinToss? column : alias[column];
    }    
}