Walker alias method implemented in Erlang

walker_alias.erl

%%% @author Sergey Prokhorov <me@seriyps.ru>
%%% @copyright (C) 2013, Sergey Prokhorov
%%% @doc
%%% @license Apache License Version 2.0
%%%
%%% Walker alias method - efficient random selection with defined probabilities.
%%% <http://en.wikipedia.org/wiki/Alias_method>
%%%
%%% > ProbabilityValueList = [{10, a}, {20, b}, {30, c}, {40, d}],
%%% > WalkerVectors = walker_alias:build(ProbabilityValueList),
%%% > RandomValue = walker_alias:choice(WalkerVectors).
%%%
%%% Probability values may be any integers
%%% RandomValue will contain 'd' in 40% cases, 'c' in 30% and so on.
%%% WalkerVectors is reusable struct - you can call choice/1 on it for many
%%% times.
%%%
%%% Ported from Python implementation:
%%% <http://code.activestate.com/recipes/576564-walkers-alias-method-for-random-objects-with-diffe/>
%%% Other implementations:
%%% Common Lisp: http://prxq.wordpress.com/2006/04/23/more-on-the-alias-method/
%%% Ruby: https://github.com/cantino/walker_method
%%% @end
%%% Created : 16 May 2013 by Sergey Prokhorov <me@seriyps.ru>

-module(walker_alias).

-export([build/1, choice/1]).
%% -export([dist_test/0]).


-type walker_vectors() :: {walker_vectors,
                           non_neg_integer(),
                           array(),
                           array(),
                           array()}.


-spec build([{non_neg_integer(), any()}]) -> walker_vectors().
build(WeightedList) ->
    {Weights, Keys} = lists:unzip(WeightedList),
    build(Weights, Keys).

-spec build([non_neg_integer()], [any()]) -> walker_vectors().
build(Weights, Keys) ->
    N = length(Weights),
    Sumw = lists:sum(Weights),
    Prob = [W * N / Sumw || W <- Weights],
    {Short, Long} = split_short_long(Prob),
    {Inx, Prob2} = build_vectors(Short, Long, array:new(N), array:from_list(Prob)),
    {walker_vectors, N, array:from_list(Keys), Inx, Prob2}.

split_short_long(Prob) ->
    SplitFun = fun(V, {Sh, Lo, I}) when V < 1 ->
                       {[I | Sh], Lo, I + 1};
                  (V, {Sh, Lo, I}) when V >= 1 ->
                       {Sh, [I | Lo], I + 1}
               end,
    {Short, Long, _} = lists:foldl(SplitFun, {[], [], 0}, Prob),
    {Short, Long}.

build_vectors(S, L, Inx, Prob) when (S == []) orelse (L == []) ->
    {Inx, Prob};
build_vectors([J | Short], [K | Long], Inx, Prob) ->
    NewInx = array:set(J, K, Inx),
    ProbJ = array:get(J, Prob),
    ProbK = array:get(K, Prob),
    NewProbK = ProbK - (1 - ProbJ),
    NewProb = array:set(K, NewProbK, Prob),
    {NewShort, NewLong} = case NewProbK < 1 of
                              true -> {[K | Short], Long};
                              false -> {Short, [K | Long]}
                          end,
    build_vectors(NewShort, NewLong, NewInx, NewProb).


-spec choice(walker_vectors()) -> any().
choice({walker_vectors, N, Keys, Inx, Prob}) ->
    J = crypto:rand_uniform(0, N),
    U = crypto:rand_uniform(0, 100) / 100,
    Idx = case U =< array:get(J, Prob) of
              true -> J;
              false -> array:get(J, Inx)
          end,
    array:get(Idx, Keys).


%% dist_test() ->
%%     N = 1000,
%%     RandWeights = [C || <<C>> <= crypto:rand_bytes(N)],
%%     Iters = N * 100,
%%     Keys = lists:seq(1, N),
%%     KW = lists:zip(RandWeights, Keys),
%%     Vectors = build(KW),
%%     {Time, Res} = timer:tc(fun run_n_iterations/3, [Vectors, orddict:new(), Iters]),
%%     %% io:format("KW:~n~p~nRes:~n~p~n", [KW, Res]),
%%     io:format("Time: ~p Iters:~p~n", [Time, Iters]).


%% run_n_iterations(_, Dict, 0) ->
%%     Dict;
%% run_n_iterations(Vect, Dict, N) ->
%%     Val = choice(Vect),
%%     run_n_iterations(Vect, orddict:update_counter(Val, 1, Dict), N - 1).