Java实现 LeetCode 528 按权重随机选择（TreeMap）

528. 按权重随机选择

给定一个正整数数组 w ，其中 w[i] 代表位置 i 的权重，请写一个函数 pickIndex ，它可以随机地获取位置 i，选取位置 i 的概率与 w[i] 成正比。

说明:

1 <= w.length <= 10000

1 <= w[i] <= 10^5

pickIndex 将被调用不超过 10000 次

示例1:

输入:

[“Solution”,“pickIndex”]

[[[1]],[]]

输出: [null,0]

示例2:

输入:

[“Solution”,“pickIndex”,“pickIndex”,“pickIndex”,“pickIndex”,“pickIndex”]

[[[1,3]],[],[],[],[],[]]

输出: [null,0,1,1,1,0]

输入语法说明：

输入是两个列表：调用成员函数名和调用的参数。Solution 的构造函数有一个参数，即数组 w。pickIndex 没有参数。输入参数是一个列表，即使参数为空，也会输入一个 [] 空列表。

PS：

偷个懒

class Solution {

     int sum=0;
    private TreeMap<int[], Integer> range = new TreeMap<>(new Comparator<int[]>() {
        @Override
        public int compare(int[] o1, int[] o2) {
            // 区间内
            if (o1[0] >= o2[0] && o1[1] < o2[1]) {
                return 0;
            // 小于，左区间
            } else if (o1[1] <= o2[0]) {
                return -1;

            // 大于
            } else {
                return 1;
            }
        }
    });

    public Solution(int[] w) {
         int start;
        for(int i=0;i<w.length;i++) {
            start = sum;
            sum += w[i];
            range.put(new int[]{start, sum}, i);
        }
    }

    public int pickIndex() {
        int index = (int)(Math.random() * sum);
        if (range.get(new int[]{index, index}) == null) {
            return 0;
        }else{
            return range.get(new int[]{index, index});
        }
    }
}

/**
 * Your Solution object will be instantiated and called as such:
 * Solution obj = new Solution(w);
 * int param_1 = obj.pickIndex();
 */