Leetcode 528. Random Pick with Weight (python+cpp)

题目

解法：brutal force

把每个ind按照相应权重的数量扩充，比如[1,3]扩充成[0,0,1,1,1]，然后从中随机选取一个值返回。这样是TLE的，因为时间复杂度会是O(sum(weight))，不过这是最直观的
python代码如下：

class Solution:def __init__(self, w: List[int]):new_w = []for ind,w_ in enumerate(w):new_w += [ind]*w_self.new_w = new_wdef pickIndex(self) -> int:ind = random.randrange(0,len(self.new_w))return self.new_w[ind]

解法2：前缀和+线性搜索

class Solution {
    vector<int> prefix_sum;
public:Solution(vector<int>& w) {
    int tmp = 0;for(auto num : w){
    tmp += num;prefix_sum.push_back(tmp);}}int pickIndex() {
    float rand_num = (float) rand() / RAND_MAX;float target = rand_num * prefix_sum.back();for(int i=0;i<prefix_sum.size();i++){
    if(target < prefix_sum[i]){
    return i;}}return prefix_sum.size() - 1;}
};

解法3：前缀和加二分搜索

构造一个数组，每个位置的值代表这个位置之前所有权重的和，长生随机值二分搜索所在的区间并返回相应下标即可。时间复杂度为O(len(w))，这个复杂度来源于构造前缀和数组，而二分搜索复杂度仅为O(log(len(w)))。
python代码如下：

class Solution:def __init__(self, w: List[int]):self.prefix_sums = []prefix_sum = 0for weight in w:prefix_sum += weightself.prefix_sums.append(prefix_sum)self.total_sum = prefix_sumdef pickIndex(self) -> int:target = self.total_sum * random.random()l,h = 0,len(self.prefix_sums)# binary searchwhile l<h:mid = l+(h-l)//2if self.prefix_sums[mid]>target:h = midelse:l = mid+1return l

C++代码如下：
这边前缀和数组的构造利用了C++的built in function: prefix_sum. 二分搜索也用了lower_bound

class Solution {
    vector<int> prefix_sum;
public:Solution(vector<int>& w) {
    int tmp = 0;for(auto num : w){
    tmp += num;prefix_sum.push_back(tmp);}}int pickIndex() {
    float rand_num = (float) rand() / RAND_MAX;float target = rand_num * prefix_sum.back();// binary searchint l=0, r = prefix_sum.size() - 1;int mid;while(l+1<r){
    mid = (r-l)/2 + l;if(prefix_sum[mid]>target){
    r = mid;}else{
    l = mid;}}return prefix_sum[l]>target ? l : r;}
};

参考链接：https://github.com/changgyhub/notes/blob/master/leetcode/leetcode-601-650.md
https://leetcode.com/problems/random-pick-with-weight/solution/