面试算法71：按权重生成随机数

题目

输入一个正整数数组w，数组中的每个数字w[i]表示下标i的权重，请实现一个函数pickIndex根据权重比例随机选择一个下标。例如，如果权重数组w为[1，2，3，4]，那么函数pickIndex将有10%的概率选择0、20%的概率选择1、30%的概率选择2、40%的概率选择3。

分析

可以创建另一个和权重数组的长度一样的数组sums，新数组的第i个数值sums[i]是权重数组中前i个数字之和。有了这个数组sums就能很方便地根据等概率随机生成的数字p按照权重比例选择下标。例如，累加权重数组[1，2，3，4]中的权重得到的数组sums为[1，3，6，10]。有了这个累加权重的数组之后，如果0到9之间的随机数p＜1，那么选择0；如果1≤p＜3，那么选择1；如果3≤p＜6，那么选择2；如果6≤p＜10，那么选择3。

解

public class Solution {
    private int[] sums;
    private int total;

    public Solution(int[] w) {
        sums = new int[w.length];
        for (int i = 0; i < w.length; i++) {
            total += w[i];
            sums[i] = total;
        }
    }

    public static void main(String[] args) {
        int[] w = {1,2,3,4};
        Solution solution = new Solution(w);
        System.out.println(solution.pickIndex());
    }

    public int pickIndex() {
        Random random = new Random();
        int p = random.nextInt(total);

        int left = 0;
        int right = sums.length;
        while (left <= right) {
            int mid = (left + right) / 2;
            if (sums[mid] > p) {
                if (mid == 0 || (sums[mid - 1] <= p)) {
                    return mid;
                }

                right = mid - 1;
            }
            else {
                left = mid + 1;
            }
        }

        return -1;
    }
}