528. Random Pick with Weight

Problem

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. 1 <= w.length <= 10000
2. 1 <= w[i] <= 10^5
3. pickIndex will be called at most 10000 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.

Problem Summary

Given a positive integer array w, where w[i] represents the weight of position i, write a function pickIndex that can randomly get position i, where the probability of selecting position i is proportional to w[i].

Note:

1. 1 <= w.length <= 10000
2. 1 <= w[i] <= 10^5
3. pickIndex will be called no more than 10000 times

Explanation of input syntax:

The input consists of two lists: the names of the member functions called and the parameters for the calls. The constructor of Solution has one parameter, which is the array w. pickIndex has no parameters. The input parameters are a list; even if the parameters are empty, an empty list [] will be input.

Solution Approach

• Given an array, where each element value represents the weight value of that index, pickIndex() randomly picks a position i, and the probability of this position appearing is proportional to that element value.
• Since weights are involved, this problem can first be handled using prefix sums for the weights. Randomly choose an integer x in the interval [0,prefixSum). The index i is the smallest index satisfying the condition x< prefixSum[i]; finding this index i is the final answer. Use binary search to find the index i. For some index i, all integers v satisfying prefixSum[i] - w[i] ≤ v < prefixSum[i] map to this index. Therefore, all indices are proportional to their index weights.
• Time complexity: the preprocessing time complexity is O(n), and the time complexity of pickIndex() is O(log n). Space complexity is O(n).

Code

package leetcode

import (
	"math/rand"
)

// Solution528 define
type Solution528 struct {
	prefixSum []int
}

// Constructor528 define
func Constructor528(w []int) Solution528 {
	prefixSum := make([]int, len(w))
	for i, e := range w {
		if i == 0 {
			prefixSum[i] = e
			continue
		}
		prefixSum[i] = prefixSum[i-1] + e
	}
	return Solution528{prefixSum: prefixSum}
}

// PickIndex define
func (so *Solution528) PickIndex() int {
	n := rand.Intn(so.prefixSum[len(so.prefixSum)-1]) + 1
	low, high := 0, len(so.prefixSum)-1
	for low < high {
		mid := low + (high-low)>>1
		if so.prefixSum[mid] == n {
			return mid
		} else if so.prefixSum[mid] < n {
			low = mid + 1
		} else {
			high = mid
		}
	}
	return low
}

/**
 * Your Solution object will be instantiated and called as such:
 * obj := Constructor(w);
 * param_1 := obj.PickIndex();
 */