package number;

import java.util.HashMap;

public class RandomDistribution {
    /*
     * Design and implement a class to generate random numbers in an arbitrary
     * probability distribution given by an array of integer weights, i.e. for
     * int[] w return a number, n, from 0 to w.length - 1 with probability w[n]
     * / sum(w). Using an existing random number generator with a uniform
     * distribution is permitted. Example distribution: w = 1, 2, 3, 2, 1
     * Example probabilities: w / sum = 1/9, 2/9, 1/3, 2/9, 1/9 Example results:
     * n = 0, 1, 2, 3, 4
     */
    HashMap<Integer, Integer> alias = new HashMap<Integer, Integer>();
    int sum = 0;

    public RandomDistribution(int[] weights) {
        for (int w : weights) {
            sum += w;
        }
        int slot = 0;
        int partial = weights[slot];
        for (int i = 0; i < sum; i++) {
            if (i >= partial) {
                slot += 1;
                partial += weights[slot];
            }
            alias.put(i, slot);
        }
    }

    public int next() {
        return alias.get((int) (Math.random() * (sum - 1)));
    }

    /**
     * @param args
     */
    public static void main(String[] args) {
        // TODO Auto-generated method stub
        int[] dis = new int[]{1, 2, 3, 2, 1};
        RandomDistribution rd = new RandomDistribution(dis);
        System.out.println(rd.next());
    }

}
