LeetCode 528 Random Pick with Weight

    
        标签：
        
            数组类题目
        
            前缀和
        
            二分查找
        
            LeetCode
        
    发布于：2022-02-13 15:39:15
    编辑于：2022-02-13 17:22:17
    浏览量：903

概述

https://leetcode.com/problems/random-pick-with-weight/

暴力法

每个索引按照其权重扩增那么多份，然后直接取样就好。

前缀和 + 二分查找

实际上我们没有必要真的搞个数组来，我们只要能给定一个索引，返回对应的值不就行了么。

对于这种范围问题很自然想到利用前缀和来快速查找范围。

对于前缀和数组也不要 O(n) 遍历，二分查找即可。

这里的二分查找是寻找第一个不小于目标值的值，所以 while 里条件为 l < r，且 sums[m] >= t 时将 r 设置为 m，而非 m - 1，因为此时 r 时的值可能就是答案，设置为 m - 1 的话我们就把答案搞出去了。

要注意的是，这里 int t = rand() % sums.back() + 1; 必须要 + 1，否则只能通过部分用例。为什么呢？这里的 t 代表的是从暴力法中描述的那个扩增好的数组中取样，其尺寸不就是 sums.back() 么？

这是因为，t 代表的不是索引，而是 sum！如果我们真的是用那个扩增好的数组，那 int t = rand() % sums.back() 并没有问题。但是 sum 的话，不加 1 的话我们就多了一个 0 的 sum，少了一个全部都加起来的 sum。

举个例子，[1, 2, 3, 4]，不加 1 的话，我们取 0~9，会发现 1 被命中了两次（0 和 1），而 4 少了一次。

class Solution {
public:
    vector<int> sums;
    Solution(vector<int>& w) {
        sums.push_back(w[0]);
        for (int i = 1; i < w.size(); i ++) {
            sums.push_back(w[i] + sums[i - 1]);
        }
    }
    
    int pickIndex() {
        int t = rand() % sums.back() + 1;  // Why +1? 
        int l = 0, r = sums.size() - 1;
        while (l < r) {  // Find the first index that not smaller then t
            int m = (r - l) / 2 + l;
            if (sums[m] == t) {
                r = m;
            } else if (sums[m] > t) {
                r = m;  // Why not r = m - 1? Cause m may be the answer
            } else {
                l = m + 1;
            }
        }
        return l;
    }
};

Links: leetcode-528-random-pick-with-weight