题目:
给定一个正整数数组 w ,其中 w[i] 代表下标 i 的权重(下标从 0 开始),请写一个函数 pickIndex ,它可以随机地获取下标 i,选取下标 i 的概率与 w[i] 成正比。
例如,对于 w = [1, 3],挑选下标 0 的概率为 1 / (1 + 3) = 0.25 (即,25%),而选取下标 1 的概率为 3 / (1 + 3) = 0.75(即,75%)。
也就是说,选取下标 i 的概率为 w[i] / sum(w) 。
分析:
若权重数组是[1,2,3,4],那么权重之和就是10,由于下标0对应的权重是1,那么选择0的概率是10%,以此类推,选择下标1,2和3的概率分别是20%,30%,40%。
可以创建另一个和权重数组的长度一样的数组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等等(后面同理)
该算法时间复杂度为O(logn)。
代码:
import java.util.Random;
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 int pinkIndex(){
Random random = new Random();
int i = random.nextInt(total);
int left = 0;
int right = sums.length-1;
while (left<=right){
int mid = (left+right)/2;
//数组中第一个大于p的数字要满足两个条件:一个是这个数字本身要大于p,另一个是如果它前面有数字那么前一个数字要小于或等于p
if (sums[mid] >i){
if (mid == 0||sums[mid-1] <= i){
return mid;
}
right = mid -1;
}else {
left = mid +1;
}
}
return -1;
}
}