528.random-pick-with-weight.go 663 B

12345678910111213141516171819202122232425262728293031323334353637
  1. type Solution struct {
  2. sum []int
  3. l int
  4. n int
  5. }
  6. func Constructor(w []int) (sol Solution) {
  7. n := len(w)
  8. sol.sum = make([]int, n+1)
  9. for i := range w {
  10. sol.sum[i+1] = sol.sum[i] + w[i]
  11. }
  12. sol.n, sol.l = sol.sum[n], n+1
  13. return
  14. }
  15. func (this *Solution) PickIndex() int {
  16. idx := rand.Intn(this.n) + 1
  17. beg, end := 1, this.l-1
  18. for beg <= end {
  19. mid := beg + (end - beg)
  20. if val := this.sum[mid]; val == idx {
  21. return mid - 1
  22. } else if val < idx {
  23. beg = mid + 1
  24. } else {
  25. end = mid - 1
  26. }
  27. }
  28. return beg - 1
  29. }
  30. /**
  31. * Your Solution object will be instantiated and called as such:
  32. * obj := Constructor(w);
  33. * param_1 := obj.PickIndex();
  34. */