O(n) time for pickIndex() function
class Solution {
private int[] prefixSums;
private int totalSum;
public Solution(int[] w) {
this.prefixSums = new int[w.length];
int prefixSum = 0;
for (int i = 0; i < w.length; i++) {
prefixSum += w[i];
this.prefixSums[i] = prefixSum;
}
this.totalSum = prefixSum;
}
public int pickIndex() {
double target = this.totalSum * Math.random();
int i = 0;
while (i < this.prefixSums.length) {
if (target < this.prefixSums[i]) {
return i;
}
i++;
}
return i - 1;
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(w);
* int param_1 = obj.pickIndex();
*/
O(logn) time for pickIndex() function
class Solution {
int[] prefixSum;
int totalSum;
public Solution(int[] w) {
this.prefixSum = new int[w.length];
prefixSum[0] = w[0];
totalSum = w[0];
for (int i = 1; i < w.length; i++) {
prefixSum[i] = prefixSum[i - 1] + w[i];
totalSum += w[i];
}
}
public int pickIndex() {
double target = this.totalSum * Math.random();
int left = 0;
int right = this.prefixSum.length;
while (left < right) {
int mid = left + (right - left) / 2;
if (target > this.prefixSum[mid]) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
}
}
/**
* Your Solution object will be instantiated and called as such:
* Solution obj = new Solution(w);
* int param_1 = obj.pickIndex();
*/