How to pick an element from an array with a probability distribution?

Question

I have a mathematical/programming question about a problem I am trying to solve. Given this simple array of integers:

int[] goals = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

Do you know how could I implement a function to pick an element from the array given a distributed probability like the one from this example?:

The probability of picking goals[0] is 50%
The probability of picking goals[1] is 30%
The probability of picking goals[2] is 12%
The probability of picking goals[3] is 2.5%
The probability of picking goals[4] is 0.85%
etcetera

The probability distribution is of my choice, hardcoded to keep things simple.

Thank you very much for your inputs!

Just use Math.random to generate a number, then use an if to check if it is above/below a certain value to get your "probability" — MFerguson, Nov 18 '21 at 12:30

score 3 · Answer 1 · answered Nov 18 '21 at 12:56

Let's say you specify you probabilities in an array:

double[] probs = {50, 30, 12, 2.5, 0.85 /*, ...*/};

You can calculate the total of the probabilities:

double totalProbs = DoubleStream.of(probs).sum();

And declare a random object:

Random random = new Random();

You can then write a method that returns a random goal:

public int randomGoal() {
    double x = random.nextDouble()*totalProbs;
    for (int i = 0; i < probs.length; ++i) {
        x -= probs[i];
        if (x <= 0) {
            return goals[i];
        }
    }
    return goals[goals.length-1];
}

asinkxcoswt · Accepted Answer · 2021-11-18T13:15:09.957

You can create an array of 100 elements, then for each number i in the array goals, put i into the new array n times where n equals the probability of i times the array size (100). Then randomly pick an element from the new array.

You can increase the array size (100) to get a more precise result. The most precise result will be when you pick the array size that makes every n a whole number.

Example:

https://jsfiddle.net/o5hmjxd2/28/

const x = [0,1,2,3,4,5,6,7,8,9];
const p = [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
const l = 100;

let y = [];


for (const i = 0; i < x.length; i++) {
  y = y.concat(new Array(l*p[i]).fill(x[i]));
}

function getRandomInt(max) {
  return Math.floor(Math.random() * max);
}

// result
console.log(y);
console.log(y[getRandomInt(y.length - 1)])

score 1 · Answer 3 · answered Nov 18 '21 at 13:42

Try this.

public class DistributedProbability {

    final NavigableMap<Double, Integer> distribution = new TreeMap<>();
    final Random random = new Random();
    public final int size;
    public final double max;

    public DistributedProbability(double... rates) {
        double sum = 0;
        int index = 0;
        for (double rate : rates) {
            sum += rate;
            this.distribution.put(sum, index++);
        }
        size = index;
        max = sum;
    }

    int next() {
        double d = random.nextDouble(max);
        Entry<Double, Integer> entry = distribution.higherEntry(d);
        return entry.getValue();
    }
}

and

public static void main(String[] args) {
    DistributedProbability d = new DistributedProbability(
        50, 30, 12, 2.5, 0.85, 4.65);

    int[] counts = new int[d.size];
    for (int i = 0; i < 100000; ++i)
        ++counts[d.next()];

    System.out.println(Arrays.toString(counts));
}

output:

[49844, 30101, 12023, 2512, 845, 4675]

How to pick an element from an array with a probability distribution?

3 Answers3