Select array element by its weight/probability

Question

I coded a function that takes an array with numbers and randomly chooses one of these numbers (Returns the index of the chosen one). The value of each number works as an weight which determines the probability to that number be chosen.

E.G.: A possible input for the function could be [1.5, 3.5, 2.0, 1.3] and the respective probabilities would be 18.07%, 42.16%, 24.09%, 15.66%. In other words, the probability to choose a number is proportional to that number.

The function:

function chooseIndex(weights){

    let weightSum = 0;
    weights.forEach(weight => {
        weightSum += weight;
    });

    let choice = Math.random()*weightSum; // [0 ~ weightSum[
    let retIndex;
    weights.some((weight, index) => {
        choice -= weight;
        if(choice < 0){
            retIndex = index;
            return true;
        }
    });
    return retIndex;
}

I don't know why, but I feel it could be better about performance.
Any suggestions will be very welcome ;D

Answer 1

More concise code

First, you can compute weightSum using a reduce() function instead. I am not sure if there is any particular performance benefit to this or not, but it does allow you to put everything on one line if you like that sort of thing.

const weightSum = weights.reduce((sum, weight) => sum + weight, 0);

Similarly, in newer versions of Javascript you can use weights.findIndex() instead of weights.some() to make your code more concise. However this will not work in Internet Explorer.

    const retIndex = weights.findIndex((weight) => {
        choice -= weight;
        return (choice < 0);
    });
    return retIndex;

Now, on to efficiency.

Efficiency - using input once

If you only plan to use each input once, then the algorithm you have is close to the best fit for the job. But the time to produce an output increases linearly with the size of the input (O(n)). We can improve that. (Thanks Roland Illig.)

Say that at the top, instead of simply taking the sum, we generate a new array with all the intermediate steps to get to the sum. E.g. for the input [1.5, 3.5, 2, 1.3] we create the array [1.5, 5, 7, 8.3]. This leads us to more or less the inverse of your current method; instead of subtracting weights from one side of the less-than comparison (i.e. from choice), we are adding them to the other side.

    const accumulatedWeights = weights.slice(0, 1);
    weights.slice(1).forEach((weight, i) => {
        accumulatedWeights[i] = weight + accumulatedWeights[i - 1];
    });
    weightSum = accumulatedWeights[accumulatedWeights.length - 1];

    let choice = Math.random() * weightSum; // [0 ~ weightSum]
    let retIndex;
    accumulatedWeights.some((accumulatedWeight, index) => {
        if(choice < accumulatedWeight){
            retIndex = index;
            return true;
        }
    });
    return retIndex;

Now, this is not yet an improvement on your current method. The key is that this new array accumulatedWeights has a useful property: we know it is sorted, ordered from least to most (assuming no negative weights, which would frankly break everything anyway). This means we can do a binary search on it. Instead of using .some() to go from one end of the array to the other, we can start by checking in the middle of the array, and then "zoom in" on the half that must contain the correct value, then repeat.

Think about if there are 100 values in the input, and the random number generated corresponds to the value at index 99. Without binary search, we have to go through the whole array and do 100 operations to find the correct index; with binary search we only have to do around 7. (All 100 -> top 50 -> top 25 -> top 12 or 13 -> top 6 or 7 -> top 3 or 4 -> top 1 or 2 -> definitely top 1.) This is O(log n) time efficiency, a big improvement over O(n).

Unfortunately I don't have time to write out a full binary search function, but see this StackOverflow post.

Efficiency - using input many times

But what if you wanted to use each input more than once? E.g. you want to select ten million randomized indexes in a row based on the same weighting scheme.

Well, in that case you could follow the advice of this other StackOverflow post:

• First, build a big array where each number you want to select appears some number of times in proportion to its weight. E.g. for the weights [1.5, 3.5, 2, 1.3], you would generate an array with fifteen 0s, thirty-five 1s, twenty 2s, and thirteen 3s.
• Then, pick a random element from the big array (i.e. compute an index via Math.floor(Math.random() * big_array.length) and return the element at that index) ten million times in a row.

You pay a certain cost up front to set up the big array, but once you have it set up, the time to find another result is constant (O(1)), no matter how long the original input was.

Whether or not this is more efficient in practice depends entirely on your use case.