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

蔚落 2023-07-18 04:48 27阅读 0赞

## 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(); */