LeetCode 528 Random Pick with Weight

Tag: 数组类题目 前缀和 二分查找 LeetCode Posted on 2022-02-13 15:39:15 Edited on 2022-02-13 17:22:17 Views: 160

概述

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;
    }
};

未经允许,禁止转载,本文源站链接:https://iamazing.cn/