Random pick with weight

Array Math Binary Search Prefix Sum Randomized

48.9% acceptance

Feb 19, 2026

2267

1054

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%).

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

* impl Solution {

 *     fn new(w: Vec<i32>) -> Self {

 *     }

 *     fn pick_index(&self) -> i32 {

 *     }
 * }
 */

/**
 * Your Solution object will be instantiated and called as such:
 * let obj = Solution::new(w);
 * let ret_1: i32 = obj.pick_index();
 */

use rand::Rng;
struct Solution {
  prefix: Vec<i32>,
  total: i32,
}
impl Solution {
  #[allow(dead_code)]
  fn new(w: Vec<i32>) -> Self {
    let mut prefix = Vec::with_capacity(w.len());
    let mut sum = 0i32;
    for x in w { sum += x; prefix.push(sum); }
    Solution { total: sum, prefix }
  }
  fn pick_index(&self) -> i32 {
    let mut rng = rand::rng();
    let target = rng.random_range(1..=self.total);
    let mut lo = 0i32;
    let mut hi = self.prefix.len() as i32 - 1;
    while lo < hi {
      let mid = (lo + hi) / 2;
      if self.prefix[mid as usize] < target { lo = mid + 1; } else { hi = mid; }
    }
    lo
  }
}