Generate one of 3 numbers according to given probabilities

Problem

You are given a function rand(a, b) which generates equiprobable random numbers between [a, b] inclusive. Generate one of three numbers x, y, or z with probabilities P(x), P(y), and P(z) such that P(x) + P(y) + P(z) = 1 using the given rand(a, b) function.

Examples

Example 1

Input: x = 1, y = 2, z = 3, P(x) = 0.2, P(y) = 0.5, P(z) = 0.3
Output: 2 (or 1 or 3, according to the given probabilities)
Explanation: Each number is returned with its specified probability.

Example 2

Input: x = 10, y = 20, z = 30, P(x) = 0.1, P(y) = 0.7, P(z) = 0.2
Output: 20 (or 10 or 30, according to the given probabilities)
Explanation: Each number is returned with its specified probability.

Solution

Method 1 – Cumulative Probability and Uniform Random

Intuition

To generate one of three numbers with specified probabilities, we can use a uniform random number and map intervals of the range [0, 1) to each number according to its probability. This ensures each number is selected with the correct probability.

Approach

1. Compute the cumulative probabilities for x, y, and z.
2. Generate a uniform random number r in [0, 1) using the given rand(a, b) function.
3. If r < P(x), return x.
4. Else if r < P(x) + P(y), return y.
5. Else, return z.

Code

C++

class Solution {
public:
  int generate(int x, int y, int z, double px, double py, double pz) {
    double r = (double)rand() / RAND_MAX;
    if (r < px) return x;
    else if (r < px + py) return y;
    else return z;
  }
};

Go

func Generate(x, y, z int, px, py, pz float64) int {
  r := rand.Float64()
  if r < px {
    return x
  } else if r < px+py {
    return y
  } else {
    return z
  }
}

Java

class Solution {
  public int generate(int x, int y, int z, double px, double py, double pz) {
    double r = Math.random();
    if (r < px) return x;
    else if (r < px + py) return y;
    else return z;
  }
}

Kotlin

class Solution {
  fun generate(x: Int, y: Int, z: Int, px: Double, py: Double, pz: Double): Int {
    val r = Math.random()
    return when {
      r < px -> x
      r < px + py -> y
      else -> z
    }
  }
}

Python

class Solution:
  def generate(self, x: int, y: int, z: int, px: float, py: float, pz: float) -> int:
    r = random.random()
    if r < px:
      return x
    elif r < px + py:
      return y
    else:
      return z

Rust

use rand::Rng;
impl Solution {
  pub fn generate(x: i32, y: i32, z: i32, px: f64, py: f64, pz: f64) -> i32 {
    let r = rand::thread_rng().gen::<f64>();
    if r < px {
      x
    } else if r < px + py {
      y
    } else {
      z
    }
  }
}

TypeScript

class Solution {
  generate(x: number, y: number, z: number, px: number, py: number, pz: number): number {
    const r = Math.random();
    if (r < px) return x;
    else if (r < px + py) return y;
    else return z;
  }
}

Complexity

• ⏰ Time complexity: O(1) — Only a few arithmetic operations and a random number generation.
• 🧺 Space complexity: O(1) — Only a few variables are used.