Description
You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith index.
You need to implement the function pickIndex(), which randomly picks an index in the range [0, w.length - 1] (inclusive) and returns it. The probability of picking an index i is w[i] / sum(w).
- For example, if
w = [1, 3], the probability of picking index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%), and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).
Example 1:
Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]
Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.
Example 2:
Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]
Explanation
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // return 1. It is returning the second element (index = 1) that has a probability of 3/4.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It is returning the first element (index = 0) that has a probability of 1/4.
Since this is a randomization problem, multiple answers are allowed.
All of the following outputs can be considered correct:
[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]
......
and so on.
Constraints:
1 <= w.length <= 1041 <= w[i] <= 105pickIndex will be called at most 104 times.
Solutions
Solution 1
PythonJavaC++GoRustJavaScript
class Solution:
def __init__(self, w: List[int]):
self.s = [0]
for c in w:
self.s.append(self.s[-1] + c)
def pickIndex(self) -> int:
x = random.randint(1, self.s[-1])
left, right = 1, len(self.s) - 1
while left < right:
mid = (left + right) >> 1
if self.s[mid] >= x:
right = mid
else:
left = mid + 1
return left - 1
# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()(code-box)
class Solution {
private int[] s;
private Random random = new Random();
public Solution(int[] w) {
int n = w.length;
s = new int[n + 1];
for (int i = 0; i < n; ++i) {
s[i + 1] = s[i] + w[i];
}
}
public int pickIndex() {
int x = 1 + random.nextInt(s[s.length - 1]);
int left = 1, right = s.length - 1;
while (left < right) {
int mid = (left + right) >> 1;
if (s[mid] >= x) {
right = mid;
} else {
left = mid + 1;
}
}
return left - 1;
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(w);
* int param_1 = obj.pickIndex();
*/(code-box)
class Solution {
public:
vector<int> s;
Solution(vector<int>& w) {
int n = w.size();
s.resize(n + 1);
for (int i = 0; i < n; ++i) s[i + 1] = s[i] + w[i];
}
int pickIndex() {
int n = s.size();
int x = 1 + rand() % s[n - 1];
int left = 1, right = n - 1;
while (left < right) {
int mid = left + right >> 1;
if (s[mid] >= x)
right = mid;
else
left = mid + 1;
}
return left - 1;
}
};
/**
* Your Solution object will be instantiated and called as such:
* Solution* obj = new Solution(w);
* int param_1 = obj->pickIndex();
*/(code-box)
type Solution struct {
s []int
}
func Constructor(w []int) Solution {
n := len(w)
s := make([]int, n+1)
for i := 0; i < n; i++ {
s[i+1] = s[i] + w[i]
}
return Solution{s}
}
func (this *Solution) PickIndex() int {
n := len(this.s)
x := 1 + rand.Intn(this.s[n-1])
left, right := 1, n-1
for left < right {
mid := (left + right) >> 1
if this.s[mid] >= x {
right = mid
} else {
left = mid + 1
}
}
return left - 1
}
/**
* Your Solution object will be instantiated and called as such:
* obj := Constructor(w);
* param_1 := obj.PickIndex();
*/(code-box)
use rand::{thread_rng, Rng};
struct Solution {
sum: Vec<i32>,
}
/**
* `&self` means the method takes an immutable reference.
* If you need a mutable reference, change it to `&mut self` instead.
*/
impl Solution {
fn new(w: Vec<i32>) -> Self {
let n = w.len();
let mut sum = vec![0; n + 1];
for i in 1..=n {
sum[i] = sum[i - 1] + w[i - 1];
}
Self { sum }
}
fn pick_index(&self) -> i32 {
let x = thread_rng().gen_range(1, self.sum.last().unwrap() + 1);
let (mut left, mut right) = (1, self.sum.len() - 1);
while left < right {
let mid = (left + right) >> 1;
if self.sum[mid] < x {
left = mid + 1;
} else {
right = mid;
}
}
(left - 1) as i32
}
}(code-box)
/**
* @param {number[]} w
*/
var Solution = function (w) {
const n = w.length;
this.s = new Array(n + 1).fill(0);
for (let i = 0; i < n; ++i) {
this.s[i + 1] = this.s[i] + w[i];
}
};
/**
* @return {number}
*/
Solution.prototype.pickIndex = function () {
const n = this.s.length;
const x = 1 + Math.floor(Math.random() * this.s[n - 1]);
let left = 1,
right = n - 1;
while (left < right) {
const mid = (left + right) >> 1;
if (this.s[mid] >= x) {
right = mid;
} else {
left = mid + 1;
}
}
return left - 1;
};
/**
* Your Solution object will be instantiated and called as such:
* var obj = new Solution(w)
* var param_1 = obj.pickIndex()
*/(code-box)
