Leetcode_按权重随机选择

题目

给定一个正整数数组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)。

 

题解

思路：对w的元素值进行前缀求和得到sum，利用rand()函数生成随机值并对sum取模以选取对应的下标。

代码：

vector<int> num;
int sum = 0;
Solution(vector<int>& w) {
   int i = 0;
    while(i < w.size()){
        sum += w[i];
        num.push_back(sum);
        i++;
    }
}

int pickIndex() {
    int i;
    int pick = rand() % sum;
    for(i = 0; i < num.size(); i++){
        if(pick < num[i])
            break;
    }
    return i;
}

逻辑上复杂度并不高，但还是跑得很慢…
 

参考大佬调用库函数的代码（思路一致）：

vector<int> sum;
Solution(vector<int>& w):sum(std::move(w)) {
    partial_sum(sum.begin(),sum.end(),sum.begin());
}

int pickIndex() {
    int pos =(rand()%sum.back())+1;
    return lower_bound(sum.begin(),sum.end(),pos) - sum.begin();
}

求w的前缀和利用了partial_sum()函数（定义在<numeric>中），不单独存储总和，而是直接通过sum最后一个元素得到。用lower_bound()函数（定义在<algorithm>中）代替for循环的遍历。

particial_sum()函数的介绍
lower_bound()函数的介绍