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

1
2
3
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

1
2
3
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

1
2
3
4
5
6
7
8
9
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;
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
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
  }
}
1
2
3
4
5
6
7
8
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;
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
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
    }
  }
}
1
2
3
4
5
6
7
8
9
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
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
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
    }
  }
}
1
2
3
4
5
6
7
8
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.