LeetCode 528. Random Pick with Weight
原题链接在这里:https://leetcode.com/problems/random-pick-with-weight/
题目:
Given an array w
of positive integers, where w[i]
describes the weight of index i
, write a function pickIndex
which randomly picks an index in proportion to its weight.
Note:
1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex
will be called at most10000
times.
Example 1:
Input:
["Solution","pickIndex"]
[[[1]],[]]
Output: [null,0]
Example 2:
Input:
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output: [null,0,1,1,1,0]
Explanation of Input Syntax:
The input is two lists: the subroutines called and their arguments. Solution
's constructor has one argument, the array w
. pickIndex
has no arguments. Arguments are always wrapped with a list, even if there aren't any.
题解:
The given input array means for current index i, the weight is w[i].
e.g. w = [20, 20, 60]. The weight of choosing 0 is 20, the weight of choosing 1 is 20, the weight of choosing 2 is 60, totally it is 100.
In order to choose by weight, we need to accumlat the total weight, and pick a random number within [1, total weight].
sum = [20, 40, 100]
And binary search where this weight would land.
If sum[mid] == pick, then simply return mid.
Else if sum[mid] < pick, say pick = 50, sum[mid] = 40, then it can't be index 1, since 1 is from 20 to 40. l = mid + 1.
Else if sum[mid] > pick, say pick = 30, sum[mid] = 40, then it could still be index 1, since 1 is from 20 to 40. r = mid.
When using above transion, while loop condision is l < r, it can't be l <= r. Otherwise, it would get out of loop.
Time Complexity: pickIndex, O(logw.length).
Space: O(1). It doesn't have extra array, it is changing input w array.
AC Java:
1 class Solution { 2 int [] sum; 3 Random rand; 4 5 public Solution(int[] w) { 6 for(int i = 1; i<w.length; i++){ 7 w[i] += w[i-1]; 8 } 9 10 this.sum = w; 11 this.rand = new Random(); 12 } 13 14 public int pickIndex() { 15 int len = sum.length; 16 int n = sum[len-1]; 17 int pick = rand.nextInt(n) + 1; 18 19 int l = 0; 20 int r = len-1; 21 while(l < r){ 22 int mid = l + (r - l) / 2; 23 if(sum[mid] == pick){ 24 return mid; 25 }else if(sum[mid] < pick){ 26 l = mid + 1; 27 }else{ 28 r = mid; 29 } 30 } 31 32 return l; 33 } 34 } 35 36 /** 37 * Your Solution object will be instantiated and called as such: 38 * Solution obj = new Solution(w); 39 * int param_1 = obj.pickIndex(); 40 */