题目
给定一个正整数数组 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:
输入:
["Solution","pickIndex"]
[[[1]],[]]
输出:
[null,0]
解释:
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0,因为数组中只有一个元素,所以唯一的选择是返回下标 0。
示例 2:
输入:
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
输出:
[null,1,1,1,1,0]
解释:
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1,返回下标 1,返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0,返回下标 0,返回该下标概率为 1/4 。
由于这是一个随机问题,允许多个答案,因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。
官方题解
前缀和与二分搜索
想法
如果我们从 半开区间 [0,tot) 中随机选择一个整数会发生什么?
是否有办法将每一个可能的整数映射到 w 中一个下标,使得每个下标映射的数目与下标的权重对应呢?
是否有办法使用少于 O(tot) 的空间呢?
class Solution {
List<Integer> psum = new ArrayList<>();
int tot = 0;
Random rand = new Random();
public Solution(int[] w) {
for (int x : w) {
tot += x;
psum.add(tot);
}
}
public int pickIndex() {
int targ = rand.nextInt(tot);
int lo = 0;
int hi = psum.size() - 1;
while (lo != hi) {
int mid = (lo + hi) / 2;
if (targ >= psum.get(mid)) lo = mid + 1;
else hi = mid;
}
return lo;
}
}
大佬解法
https://leetcode-cn.com/problems/random-pick-with-weight/solution/lun-pan-du-suan-fa-by-joyboy/
一、题目解析
题目要求设计这样一个数据结构:初始给定一个数组w,w[i]表示下标i的权重,有一个函数pickIndex,其每次返回一个下标。要求当pickIndex执行次数足够多时,下标i在统计意义上的出现概率满足(下标i的权重/总权重)。
二、解题思路
这涉及了一个比较很容易理解的算法————LPD算法。假设w数组为[1, 3],即下标0的权重为1,下标1的权重为3,那么当执行pickIndexndex的次数足够多时,下标0的出现次数:下标1的出现次数=1:3。我们假设有一个圆盘,把它划成4等份,其中0占1份,1占3份,当固定一个指针后旋转圆盘,圆盘停止后指针指向0的概率和指向1的概率就是1:3。
用代码实现时,我们像求前缀和那样求概率前缀和,在w数组为[1, 3]的例子中,这个前缀和为[1, 4],我们把它化成分数,得到概率前缀和为[1/4, 1]。其实也可以不用这一步,只不过分数比较容易体现概率。
现在我们随机生成一个小数randomNum,如果它落在(0, 1/4],表示选中的下标的概率小于等于1/4,也就是0。如果它落在(1/4, 1],表示选中的下标的概率大于1/4,小于等于1,也就是1。那么只要用二分查找,在概率前缀和数组(前缀和数组是递增的,也就是有序的)中找到第一个小于等于randomNum,就可以知道它落在哪个区间内。注意,如果没有找到比randomNum小的数,说明它选中的是下标0。如果第一个比randomNum的概率和的下标为0,说明选中的是下标1,以此类推。
class Solution {
private double[] prePrbblty = null;
private Random random = new Random();
public Solution(int[] w) {
double sum = 0;
for (int i = 0; i < w.length; i++) //求权重总和
sum += w[i];
prePrbblty = new double[w.length];
for (int i = 0; i < w.length; i++) //求概率前缀和
prePrbblty[i] = w[i] + (i == 0 ? 0 : prePrbblty[i - 1]);
for (int i = 0; i < w.length; i++) //化成分数
prePrbblty[i] /= sum;
}
public int pickIndex() {
double rd = random.nextDouble(); //生成随机数
//二分查找找第一个小于等于随机数的前缀和
int l = 0, r = prePrbblty.length - 1;
while (l < r) {
int mid = (l + r + 1) / 2;
if (prePrbblty[mid] <= rd) {
l = mid;
} else {
r = mid - 1;
}
}
if (prePrbblty[l] > rd) //如果没有比随机数小的,说明选中下标0
return 0;
else
return l + 1;
}
}
Java 前缀和+二分查找,逐行注释
https://leetcode-cn.com/problems/random-pick-with-weight/solution/java-qian-zhui-he-er-fen-cha-zhao-zhu-xi-v6u2/
解题思路
该题思路来自官方题解,加了一点自己的理解和注释。
这道题可以抽象成:将整个数组平铺起来,然后随即扔一个小球,掉到哪个分组里,就返回哪个分组的下标。
拿 [1, 3] 来举例子,
它们映射到概率分布上,取下标 0 的概率应该是 1/4 取下标 1 的概率应该是 3/4。
而前缀和数组也正好如此划分,前缀和数组为 [ 1,4 ]。
那么取到 [ 0,1)范围内的概率和取到下标 0 的概率是一样的,对于取到 1 和取到 [ 1,4)范围内的概率也是一样的。
所以取下标的问题就转换成了取某一个和的问题,唯一需要注意的是,使用随机数产生器取到的和 target,需要使用二分法来找到最近的一个前缀和大于 target 的的下标。
class Solution {
int[] prefix; // 前缀和数组
int sum = 0; // 总和
Random random = new Random(); // 随机数产生器
public Solution(int[] w) {
int n = w.length;
prefix = new int[n];
prefix[0] = w[0];
// 建立前缀和数组
for (int i = 1; i < n; ++i) prefix[i] = prefix[i - 1] + w[i];
// 求和
sum = prefix[n - 1];
}
public int pickIndex() {
// 先使用随机数产生器来得到需要的前缀和
int target = random.nextInt(sum);
// 然后使用二分查找来逐步搜索
int left = 0, right = prefix.length - 1;
while (left < right) {
int mid = left + (right - left) / 2;
// 如果前缀和偏小了,就需要把左指针往右移动,这里需要注意的是,
// 等于情况也不行,因为这样等于是碰到了下一个区间的起始位置
if (prefix[mid] <= target) left = mid + 1;
else right = mid;
}
return left;
}
}
}