[LeetCode]Random Pick with Weight

我们首先看如果array只有两个元素该如何实现Random Pick with Weight，假设数组为[3, 5]，那么显然我们应该实现3/8的概率选到3，5/8的概率选到5。那么显而易见的，我们应该生成[1, 8]的随机数，如果随机数落在[1, 3]我们选3，否则选5。这个思路很容易推广，如果我们计算输入数组的presum数组和总的和sum，我们只需要生成[1, sum]的随机数r，去presum里找大于等于r的最小的数对应的index即可。weight都是大于0的，presum数组是严格递增的，我们binary search即可。时间复杂度preprocess生成presum数组O(n)，n为输入数组长度，pick的复杂度为O(log n)，代码如下：

	class Solution {
	public:
	Solution(vector<int> w) {
	sum = 0;
	for(const auto& num : w)
	{
	sum += num;
	presums.push_back(sum);
	}
	}
	int pickIndex() {
	int len = presums.size(), lo = 0, hi = len - 1, r = rand() % sum + 1;
	while(lo < hi)
	{
	int mid = lo + (hi - lo) / 2;
	if(presums[mid] > r)hi = mid;
	else if(presums[mid] < r)lo = mid + 1;
	else return mid;
	}
	return lo;
	}
	private:
	vector<int> presums;
	int sum;
	};
	/**
	* Your Solution object will be instantiated and called as such:
	* Solution obj = new Solution(w);
	* int param_1 = obj.pickIndex();
	*/

view raw Random Pick with Weight.cpp hosted with ❤ by GitHub