Q528. 按權重隨機選擇

題目描述

給定一個正整數陣列 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) 。

class Solution {

    public Solution(int[] w) {
      
    }

    public int pickIndex() {
        
    }
}

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

示例一

輸入：
["Solution","pickIndex"]
[[[1]],[]]
輸出：
[null,0]
解釋：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因為陣列中只有一個元素，所以唯一的選擇是返回下標 0。

示例二

輸入：
["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]
......
諸若此類。

分析

按題目描述，帶權隨機，之前我的思路是先將陣列w的和計算出來，然後新建一個double[] chance陣列，chance[i] = chance[i-1]+nums[i]/sum; 但是這樣在寫二分法的時候處理邊界不是很好寫，於是我們利用字首和，計算出preSum,然後在[0,sum]中建立一個隨機數，然後使用二分法找到對應的下標（這裡是右側二分法！！）。

程式碼

public class Solution {
    int[] sum;

    public Solution(int[] w) {
        int n = w.length;
        sum = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            sum[i] = sum[i - 1] + w[i - 1];
        }
    }

    public int pickIndex() {
        int n = sum.length;
        int t = (int) (Math.random() * sum[n - 1]) + 1;
        int l = 1, r = n - 1;
        while (l < r) {
            // >>為移位符，以下寫法等價於 (l+r)/2
            int mid = l + r >> 1;
            if (sum[mid] >= t) {
                r = mid;
            } else {
                l = mid + 1;
            }
        }
        return r - 1;
    }
}