Search code examples
iosswiftcoding-style

Generate random numbers with a given distribution


Check out this question:

Swift probability of random number being selected?

The top answer suggests to use a switch statement, which does the job. However, if I have a very large number of cases to consider, the code looks very inelegant; I have a giant switch statement with very similar code in each case repeated over and over again.

Is there a nicer, cleaner way to pick a random number with a certain probability when you have a large number of probabilities to consider? (like ~30)


Solution

  • This is a Swift implementation strongly influenced by the various answers to Generate random numbers with a given (numerical) distribution.

    For Swift 4.2/Xcode 10 and later (explanations inline):

    func randomNumber(probabilities: [Double]) -> Int {
    
        // Sum of all probabilities (so that we don't have to require that the sum is 1.0):
        let sum = probabilities.reduce(0, +)
        // Random number in the range 0.0 <= rnd < sum :
        let rnd = Double.random(in: 0.0 ..< sum)
        // Find the first interval of accumulated probabilities into which `rnd` falls:
        var accum = 0.0
        for (i, p) in probabilities.enumerated() {
            accum += p
            if rnd < accum {
                return i
            }
        }
        // This point might be reached due to floating point inaccuracies:
        return (probabilities.count - 1)
    }
    

    Examples:

    let x = randomNumber(probabilities: [0.2, 0.3, 0.5])
    

    returns 0 with probability 0.2, 1 with probability 0.3, and 2 with probability 0.5.

    let x = randomNumber(probabilities: [1.0, 2.0])
    

    return 0 with probability 1/3 and 1 with probability 2/3.


    For Swift 3/Xcode 8:

    func randomNumber(probabilities: [Double]) -> Int {
    
        // Sum of all probabilities (so that we don't have to require that the sum is 1.0):
        let sum = probabilities.reduce(0, +)
        // Random number in the range 0.0 <= rnd < sum :
        let rnd = sum * Double(arc4random_uniform(UInt32.max)) / Double(UInt32.max)
        // Find the first interval of accumulated probabilities into which `rnd` falls:
        var accum = 0.0
        for (i, p) in probabilities.enumerated() {
            accum += p
            if rnd < accum {
                return i
            }
        }
        // This point might be reached due to floating point inaccuracies:
        return (probabilities.count - 1)
    }
    

    For Swift 2/Xcode 7:

    func randomNumber(probabilities probabilities: [Double]) -> Int {
    
        // Sum of all probabilities (so that we don't have to require that the sum is 1.0):
        let sum = probabilities.reduce(0, combine: +)
        // Random number in the range 0.0 <= rnd < sum :
        let rnd = sum * Double(arc4random_uniform(UInt32.max)) / Double(UInt32.max)
        // Find the first interval of accumulated probabilities into which `rnd` falls:
        var accum = 0.0
        for (i, p) in probabilities.enumerate() {
            accum += p
            if rnd < accum {
                return i
            }
        }
        // This point might be reached due to floating point inaccuracies:
        return (probabilities.count - 1)
    }