Sample index from probability distribution

You are given a discrete probability distribution for an $M$-sided die as an array p[0..M-1], where each p[i] is non-negative.

Design a function sample(p) -> int that returns a random index i such that:

• $\Pr(sample(p)=i) = p[i] / \sum_j p[j]$

Requirements

• The function will be called many times with the same p .
• Aim for $O(\log M)$ time per sample after preprocessing.

Follow-up

What if the probabilities do not sum to 1? Explain how your approach handles this.

Clarifications

• Assume you have access to a random number generator that can produce a uniform float in $[0,1)$ or a uniform integer in a range.
• If $\sum_j p[j] = 0$ , you may define behavior (e.g., throw an error).