Problem

You are given n numbers as well as n probabilities that sum up to 1. Write a function to generate one of the numbers with its corresponding probability.

Examples

Example 1

Input: numbers = [1, 2, 3, 4], probabilities = [0.1, 0.5, 0.2, 0.2]
Output: Returns 1 10% of the time, 2 50% of the time, and 3 or 4 20% of the time.
Explanation: The function should utilize the given probabilities to return the numbers based on their probabilities.

Solution

Method 1 - Random algorithm

The task requires generating a number based on pre-defined probabilities. To achieve this, we can use the concept of cumulative probabilities. We transform the given probabilities into cumulative probabilities and use a uniformly generated random number to determine the output number according to these cumulative probabilities.

Approach

  1. Preprocess Probabilities: Compute cumulative probabilities from the given probabilities.
  2. Generate Random Number: Generate a uniformly random number between 0 and 1.
  3. Map the Random Number to Output: Use the generated random number to determine the corresponding number based on cumulative probabilities.

Code

Java
public class Solution {
    private int[] numbers;
    private double[] cumulativeProbabilities;
    private Random random;

    public void init(int[] numbers, double[] probabilities) {
        this.numbers = numbers;
        this.cumulativeProbabilities = new double[probabilities.length];
        this.random = new Random();

        cumulativeProbabilities[0] = probabilities[0];
        for (int i = 1; i < probabilities.length; i++) {
            cumulativeProbabilities[i] = cumulativeProbabilities[i - 1] + probabilities[i];
        }
    }

    public int generate() {
        double r = random.nextDouble();
        for (int i = 0; i < cumulativeProbabilities.length; i++) {
            if (r < cumulativeProbabilities[i]) {
                return numbers[i];
            }
        }
        return numbers[numbers.length - 1]; // In case of precision issues with double comparisons.
    }
}
Python
class Solution:
    def __init__(self):
        self.numbers = []
        self.cumulative_probabilities = []
        
    def init(self, numbers: List[int], probabilities: List[float]) -> None:
        self.numbers = numbers
        self.cumulative_probabilities = [0] * len(probabilities)
        self.cumulative_probabilities[0] = probabilities[0]
        for i in range(1, len(probabilities)):
            self.cumulative_probabilities[i] = self.cumulative_probabilities[i - 1] + probabilities[i]

    def generate(self) -> int:
        r = random.random()
        for i, cp in enumerate(self.cumulative_probabilities):
            if r < cp:
                return self.numbers[i]
        return self.numbers[-1]  # In case of precision issues with float comparisons.

Complexity

  • ⏰ Time complexity: 
    • Initialization (n numbers): O(n) for preprocessing cumulative probabilities.
    • Generation: O(1) for generating the random number and determining the output.
  • 🧺 Space complexity: O(n) for storing cumulative probabilities.