← 返回 capitalone 的题目列表ML Coding: Top-p (Nucleus) Sampling
类型：online_judge
ML Coding: Top-p (Nucleus) Sampling
Implement top-p / nucleus sampling for decoding.
Given a probability distribution p (unsorted) and a threshold p_thresh (0 < p_thresh <= 1):
Sort tokens by probability descending.
Select the smallest set S such that sum_{i in S} p_i >= p_thresh.
Renormalize probabilities within S.
Given a random number u (0 <= u < 1), perform deterministic sampling over S by locating u in the CDF.
Output the sampled token’s original index (0-based).
Input (stdin)
n
p[0..n-1] (sum to 1)
p_thresh
u
Output
sampled token index
Constraints
1 <= n <= 200000
If probabilities tie, break ties by smaller original index first.
Example
See test cases below.
Example
Input
5
0.4 0.2 0.15 0.15 0.1
0.7
0.0
Output
0