Problem
Language models pick the next token from a probability distribution. Two settings shape it:
- Top-k: keep only the k tokens with the highest logits (if logits tie, the lower index ranks higher).
- Temperature T: divide the kept logits by T before softmax. Low T makes the choice more predictable; high T makes it more random.
To make sampling testable, the random number is given:
- Keep the top k tokens, then list them by their original index (ascending).
- Compute softmax(logit / T) over the kept tokens (subtract the largest before taking e^x).
- Walk through the list adding probabilities; return the index of the first token where the running total becomes greater than
random_value. If none does (rounding), return the last kept token.
With k = 1 this is greedy decoding: always the most likely token.
Implement sample_token(logits, temperature, k, random_value).
Input
logits: One score per token in the vocabulary.
temperature: T, greater than 0.
k: How many top tokens to keep.
random_value: A random number in [0, 1).
Output
Return the index of the chosen token.
Constraints
1 ≤ k ≤ number of logits ≤ 1,000
0 < temperature ≤ 100, 0 ≤ random_value < 1
-100 ≤ each logit ≤ 100
Examples
Input:
k = 1
logits = [1.0, 3.0, 2.0, 0.5]
temperature = 1.0
random_value = 0.99
Expected output:
expected = 1Input:
long long k = 1;
std::vector<double> logits = {1.0, 3.0, 2.0, 0.5};
double temperature = 1.0;
double random_value = 0.99;
Expected output:
long long expected = 1;Input:
long k = 1;
double[] logits = new double[] {1.0, 3.0, 2.0, 0.5};
double temperature = 1.0;
double random_value = 0.99;
Expected output:
long expected = 1;Explanation
k = 1 keeps only the highest logit, 3.0 at index 1. Its probability is 1, so it is chosen whatever the random value: greedy decoding.
Input:
k = 2
logits = [2.0, 1.0, 0.0]
temperature = 1.0
random_value = 0.8
Expected output:
expected = 1Input:
long long k = 2;
std::vector<double> logits = {2.0, 1.0, 0.0};
double temperature = 1.0;
double random_value = 0.8;
Expected output:
long long expected = 1;Input:
long k = 2;
double[] logits = new double[] {2.0, 1.0, 0.0};
double temperature = 1.0;
double random_value = 0.8;
Expected output:
long expected = 1;Explanation
Top-2 keeps indices 0 and 1 (logits 2 and 1). softmax gives about 0.731 and 0.269. The running total is 0.731 after index 0 (not above 0.8) and 1 after index 1, so index 1 is chosen.
How to approach it
Try on your own first, then open one hint at a time.
Hint 1
First choose the highest-logit tokens, breaking ties by smaller index, then restore those tokens to ascending index order.
Hint 2
Convert the kept logits into probabilities after dividing by temperature, using a numerically stable softmax.
Hint 3
When scanning cumulative probability, return the first index that crosses the random value, but fall back to the last kept token if rounding prevents a crossing.
Expected complexity
O(m log m) time, O(m) extra space
References
This challenge is adapted for CodeUtility. Follow the statement and constraints above.