July 24, 2022

528. Random Pick with Weight

528. Random Pick with Weight

O(n) time for pickIndex() function

class Solution {
    
    private int[] prefixSums;
    private int totalSum;

    public Solution(int[] w) {
        this.prefixSums = new int[w.length];
        int prefixSum = 0;
        for (int i = 0; i < w.length; i++) {
            prefixSum += w[i];
            this.prefixSums[i] = prefixSum;
        }
        this.totalSum = prefixSum;
    }
    
    public int pickIndex() {
        double target = this.totalSum * Math.random();
        int i = 0; 
        while (i < this.prefixSums.length) {
            if (target < this.prefixSums[i]) {
                return i;
            }
            i++;
        }
        return i - 1;
    }
}

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

O(logn) time for pickIndex() function

class Solution {
    
    int[] prefixSum;
    int totalSum;

    public Solution(int[] w) {
        this.prefixSum = new int[w.length];
        prefixSum[0] = w[0];
        totalSum = w[0];
        for (int i = 1; i < w.length; i++) {
            prefixSum[i] = prefixSum[i - 1] + w[i];
            totalSum += w[i];
        }
    }
    
    public int pickIndex() {
        double target = this.totalSum * Math.random();
        int left = 0;
        int right = this.prefixSum.length;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (target > this.prefixSum[mid]) {
                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();
 */
comments powered by Disqus