528. Random Pick with Weight

T: pickIndex: O(logn), Solution() : O(n)
S: O(n)

class Solution {

    Random rand = new Random();
    int[] presum;
    public Solution(int[] w) {
        presum = new int[w.length];
        presum[0] = w[0];
        for (int i = 1; i < w.length; i++) {
            presum[i] = presum[i-1] + w[i];
        }
    }
    
    // ๆ‰€ไปฅๅฆ‚ๆžœ presum ๅชๆœ‰ๅ…ฉๅ€‹ๆ•ธ
    // ๅฐฑๆœƒ็›ดๆŽฅๆ˜ฏ left right
    // presum = [1, 4]
    public int pickIndex() {
        System.out.println(Arrays.toString(presum));
        int total = presum[presum.length-1];
        int index = rand.nextInt(total)+1; // 1~total
        
        int left = 0;
        int right = presum.length-1; // 1
        
        while (left < right) { // 0 < 1
            int mid = left + (right - left)/2; // 0 + 0 = 0
            if (presum[mid] < index) { // index 3, presum[0] = 1, if index = 1, right = 0, end => ans : index0
                left = mid+1; // left = 1
            } else {
                right = mid;
            }
        }
        return left;
    }
}

/**
[1, 4]  // 5/2 = 2
1  1   4
|--|---|


[3,1,2,4]

3, 4, 6, 10

1  3  4. 6  10
|--|--|--|--|

in 1~10 find x (random number with 10)

1. cal presum
2. so we can have a weight distibution in [1~10] , 10 is total

we want [1~10], so rand.nextInt(10)+1 => (0~9)+1 = 1~10

use this rand number to get a Random Pick with Weight

T: pickIndex: O(logn), Solution() : O(n)
S: O(n)

 */

class Solution {

    Random random;
    int sum[];
    public Solution(int[] w) {
        for (int i = 1; i < w.length; i++) {
            w[i] += w[i-1];
        }
        sum = w;
        random = new Random();
    }
    
    public int pickIndex() {
        int index = random.nextInt(sum[sum.length-1]) + 1;
        int left = 0;
        int right = sum.length - 1;
        while (left < right) {
            int mid = (left + right) >>> 1;

            if (sum[mid] == index) {
                return mid;
            } else if (sum[mid] < index) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        return left;
    }
}

/**
 * Your Solution object will be instantiated and called as such:
 * Solution obj = new Solution(w);
 * int param_1 = obj.pickIndex();
 */

/**
 * @param {number[]} w
 */
var Solution = function(w) {
    
    this.sum = [w[0]];
    for(let i = 1; i < w.length; i++) {
        this.sum.push(this.sum[i-1] + w[i]);
    }
    this.size = this.sum.length-1;
};

/**
 * @return {number}
 */
Solution.prototype.pickIndex = function() {
    const {sum, size} = this;
    const index = Math.floor(Math.random() * this.sum[size]+1);
    let left = 0;
    let right = size;
    
    while(left < right){
        const mid = (left + right) >> 1;
        if(sum[mid] == index)
            return mid;
        else if(sum[mid] < index)
            left = mid + 1;
        else
            right = mid;
    }
    return left;
};

/** 
 * Your Solution object will be instantiated and called as such:
 * var obj = new Solution(w)
 * var param_1 = obj.pickIndex()
 */

Last updated