6.9

From The Algorithm Design Manual Solution Wiki
Jump to navigation Jump to search

Idea. Build once the prefix-sum array pref; to sample, draw a random integer r in [1,total] and binary-search the first prefix ≥ r. Because each integer span (pref[i-1], pref[i]] has length w[i], the chance of landing in it is exactly w[i]/total.

Preprocessing O(n)
pref[i] = w[0]+…+w[i], total = pref[-1] (store for reuse).

Sampling O(log n)
1. r = random.randint(1,total)
2. idx = lower_bound(pref, r)  /* first i with pref[i] ≥ r */
3. return idx.

Pseudocode
buildPrefix(w):
  pref = []; s = 0
  for x in w:
    s += x; pref.append(s)
  return pref

pickIndex(pref):
  r = randInt(1, pref[-1])
  lo,hi = 0,len(pref)-1
  while lo<hi:
    mid=(lo+hi)//2
    if r ≤ pref[mid]: hi = mid
    else: lo = mid+1
  return lo


Python illustration
import random, bisect, itertools

class WeightedPicker:
  def __init__(self, w):
    self.pref = list(itertools.accumulate(w))
  def pick(self):
    r = random.randint(1, self.pref[-1])
    return bisect.bisect_left(self.pref, r)


Correctness. For any index i:
P(pick = i) = |(pref[i-1], pref[i]]| / total = w[i] / total, matching the required distribution.

Complexities. Memory O(n), build O(n), each pick O(log n); with linear scan the pick step can be O(n) but needs no extra structure.


Back to Chapter 6