1. 程式人生 > >[LeetCode] Random Pick with Weight 根據權重隨機取點

[LeetCode] Random Pick with Weight 根據權重隨機取點

Given an array w of positive integers, where w[i] describes the weight of index i, write a function pickIndex which randomly picks an index in proportion to its weight.

Note:

  1. 1 <= w.length <= 10000
  2. 1 <= w[i] <= 10^5
  3. pickIndex will be called at most 10000 times.

Example 1:

Input: 
["Solution","pickIndex"]
[[[1]],[]]
Output: [null,0]

Example 2:

Input: 
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output: [null,0,1,1,1,0]

Explanation of Input Syntax:

The input is two lists: the subroutines called and their arguments. Solution's constructor has one argument, the array w

pickIndex has no arguments. Arguments are always wrapped with a list, even if there aren't any.

這道題給了一個權重陣列,讓我們根據權重來隨機取點,現在的點就不是隨機等概率的選取了,而是要根據權重的不同來區別選取。比如題目中例子2,權重為 [1, 3],表示有兩個點,權重分別為1和3,那麼就是說一個點的出現概率是四分之一,另一個出現的概率是四分之三。由於我們的rand()函式是等概率的隨機,那麼我們如何才能有權重的隨機呢,我們可以使用一個trick,由於權重是1和3,相加為4,那麼我們現在假設有4個點,然後隨機等概率取一個點,隨機到第一個點後就表示原來的第一個點,隨機到後三個點就表示原來的第二個點,這樣就可以保證有權重的隨機啦。那麼我們就可以建立權重陣列的累加和陣列,比如若權重陣列為 [1, 3, 2] 的話,那麼累加和陣列為 [1, 4, 6],整個的權重和為6,我們 rand() % 6,可以隨機出範圍 [0, 5] 內的數,隨機到 0 則為第一個點,隨機到 1,2,3 則為第二個點,隨機到 4,5 則為第三個點,所以我們隨機出一個數字x後,然後再累加和陣列中查詢第一個大於隨機數x的數字,使用二分查詢法可以找到第一個大於隨機數x的數字的座標,即為所求,參見程式碼如下:

解法一:

class Solution {
public:
    Solution(vector<int> w) {
        sum = w;
        for (int i = 1; i < w.size(); ++i) {
            sum[i] = sum[i - 1] + w[i];
        }
    }
    
    int pickIndex() {
        int x = rand() % sum.back(), left = 0, right = sum.size() - 1;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (sum[mid] <= x) left = mid + 1;
            else right = mid;
        }
        return right;
    }
    
private:
    vector<int> sum;
};

我們也可以把二分查詢法換為STL內建的upper_bound函式,參見程式碼如下:

解法二:

class Solution {
public:
    Solution(vector<int> w) {
        sum = w;
        for (int i = 1; i < w.size(); ++i) {
            sum[i] = sum[i - 1] + w[i];
        }
    }
    
    int pickIndex() {
        int x = rand() % sum.back();
        return upper_bound(sum.begin(), sum.end(), x) - sum.begin();
    }
    
private:
    vector<int> sum;
};

類似題目:

Random Pick with Blacklist

參考資料: