528. 按权重随机选择

labuladong 题解

难度中等264收藏分享切换为英文接收动态反馈

给你一个 下标从 0 开始 的正整数数组 w ，其中 w[i] 代表第 i 个下标的权重。

请你实现一个函数 pickIndex ，它可以 随机地 从范围 [0, w.length - 1] 内（含 0 和 w.length - 1）选出并返回一个下标。选取下标 i 的 概率 为 w[i] / sum(w) 。

• 例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。

 

示例 1：

输入：
["Solution","pickIndex"]
[[[1]],[]]
输出：
[null,0]
解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。

示例 2：

输入：
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
输出：
[null,1,1,1,1,0]
解释：
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。

由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。

class Solution:

    def __init__(self, w: List[int]):

        self.presum = [0]
        for n in w:
            self.presum.append(self.presum[-1]+n)
    def find(self,nums,target):
        lo = 0
        hi = len(nums) - 1
        while lo <= hi:
            mid = lo + (hi-lo)//2
            if nums[mid] < target:
                lo = mid + 1
            elif nums[mid] > target:
                hi = mid - 1
            else:
                lo = mid +1
        return lo


        
    def pickIndex(self) -> int:
        target = random.randint(0,self.presum[-1]-1)
        return self.find(self.presum,target) -1 
  

# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()

 

class Solution {
public:
    vector<int> presum;
    Solution(vector<int>& w) {
        int n = w.size();
        presum = vector<int>(n+1,0);
        for(int i = 0; i < n;i++) {
            presum[i+1] = presum[i] + w[i];
        }

    }
    int find_left_bound(int target) {
        if (presum.size()==0) return -1;
        int low = 0,high = presum.size();
        while(low < high) {
            int mid = low + (high-low)/2;
            if (presum[mid]<target) {
                low = mid +1;
            } else {
                high = mid;
            }
        }
        return low;
    }
    int pickIndex() {
        //   rand() % (b-a+1)+ a ;    就表示  a~b 之间的一个随机整数     
        int target = (rand()%(presum.back() - 1 + 1)) + 1;//生成闭区间[1,total]范围内的一个随机数  
        cout << target << endl;
        return find_left_bound(target)-1;
    }
};

/**
 * Your Solution object will be instantiated and called as such:
 * Solution* obj = new Solution(w);
 * int param_1 = obj->pickIndex();
 */