1. 程式人生 > 實用技巧 >LeetCode528. 按權重隨機選擇

LeetCode528. 按權重隨機選擇


要按照概率隨機選擇一個數,可以將陣列的值看作一個區間上的長度,比如題目給的例子,當w = [1, 3]時,我們可以假設有一個一維的區間,區間前1/4是第0個數,區間的後3/4是第1個數。

區間總長度4也就是w陣列所有數的和。

我們可以在總長度範圍(0~4)內隨機選擇一個數,假設這個數是0~1,那麼就返回0,如果這個數是1~4,那麼就返回1。
這樣就解決了按照概率隨機返回的問題。

但是怎麼判斷我們隨機選擇的數該返回什麼值呢?
這裡需要用到字首和,比如上面的例子w = [1, 3],我們可以得到字首和preSum = [0, 1, 4](計算字首和時,一般前面要預留一個0,也就是w[0] + ... + w[i]的字首和實際上是preSum[i + 1],這是有字首和的計算公式決定的)。

我們在總長度範圍內隨機取的數在區間內處於哪一個字首和的範圍內,就返回那個字首和對應的下標,比如我們取到隨機數oneRandNum = 2,那麼在字首和區間裡第一個大於等於它的字首和是下標為2(在原陣列中下標為1)的字首和4,這時我們需要返回字首和為4的那個下標2(在原陣列中下標為1),所以我們需要返回lower_bound(preSum.begin() + 1, preSum.end(), oneRandNum) - preSum.begin() - 1

有點繞hh,畫個圖就理解了~

程式碼如下:

class Solution {
public:
    vector<int> preSum;
    int n;

    Solution(vector<int>& w) {
        n = w.size();
        preSum.resize(n + 1, 0);
        for(int i = 1; i <= n; ++i) {
            preSum[i] = preSum[i - 1] + w[i - 1];
        }
    }
    
    int pickIndex() {
        int oneRandNum = rand() % preSum[n] + 1;                        // 取一個隨機數,+1是為了將隨機數對映到範圍[1, 2, ... preSum[n]]內
        return lower_bound(preSum.begin() + 1, preSum.end(), oneRandNum) - preSum.begin() - 1;     // 找到第一個大於等於oneRandNum的字首和在原陣列中對應的下標,就是答案
    }
};

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