Weighted Probability Problem in Swift

Question

I was asked a weighted probability question in a technical interview a few months ago that went something like this:

Given an input of colors and an integer "weight" value, randomly return a color based on its weight.

The sample input would be something like this:

[("Red", 20), ("Blue", 50), ("Green", 70)]

The probability that each color would be returned would be something like:

Red: 20/140

Blue: 50/140

Green: 70/140

(where 140 is the sum of the weight of the input colors)

I was thinking about this recently and decided to give it another go in Swift.

My first pass was:

func weightedColor(input: [(String, Int)]) -> String {

    var total = 0

    for tuple in input {
        total += tuple.1
    }

    let rand = Int(arc4random_uniform(UInt32(total)))

    var sum = 0
    for tuple in input {
        let range = sum + tuple.1

        if rand <= range {
            return tuple.0
        } else {
            sum += tuple.1
        }
    }

    return ""
}

Then I realized I could optimize the total weight calculation with map-reduce:

let total = UInt32(input.map { $0.1 }.reduce(0, combine: +))

My current implementation is as follows:

func weightedColor(input: [(String, Int)]) -> String {

    let total = UInt32(input.map { $0.1 }.reduce(0, combine: +))

    let rand = Int(arc4random_uniform(total))

    var sum = 0
    for tuple in input {
        let range = sum + tuple.1

        if rand <= range {
            return tuple.0
        } else {
            sum += tuple.1
        }
    }

    return ""
}

I don't like the for tuple in input loop and feel it can be optimized further. How can I improve this answer?

Answer 1

There is one error in your implementation: The test if rand <= range must be if rand < range, otherwise the returned values are not correctly distributed according to the given weights. You can check that with the input list

[("Red", 1), ("Blue", 1)]

which returns only "Red" without this fix.

Apart from that, your code looks correct, but there are some possible improvements.

• The final return "" is never reached unless you made a programming error. Instead of returning some value (in order to satisfy the compiler), replace it by

  fatalError("This should never be reached")
  

  This will stop the program execution immediately, and fatalError() has the @noreturn attribute, so the compiler won't complain about a missing return value.

• Replace for tuple in input { } by for (color, weight) in input, that allows you to get rid of the non-descriptive tuple.0 and tuple.1 in the loop body.

• Add the current weight to the sum before testing against rand, that makes the range variable obsolete.

The function does then look like this:

func weightedColor(input: [(String, Int)]) -> String {

    let total = UInt32(input.map { $0.1 }.reduce(0, combine: +))
    let rand = Int(arc4random_uniform(total))

    var sum = 0
    for (color, weight) in input {
        sum += weight
        if rand < sum {
            return color
        }
    }

    fatalError("This should never be reached")
}

The next step is to make the function generic. There is nothing which is special to colors or strings, so you can simply replace String by a generic type T:

func weightedRandomElement<T>(items: [(T, Int)]) -> T {

    let total = UInt32(items.map { $0.1 }.reduce(0, combine: +))
    let rand = Int(arc4random_uniform(total))

    var sum = 0
    for (element, weight) in items {
        sum += weight
        if rand < sum {
            return element
        }
    }

    fatalError("This should never be reached")
}

Then the same function can be used to get a random color:

let colors = [("Red", 20), ("Blue", 50), ("Green", 70)]
let color = weightedRandomElement(colors)

or a random number:

let numbers = [(3.1415, 20), (2.71827, 50)]
let number = weightedRandomElement(numbers)

or anything else.

---

Update: As @hashemi correctly said, the fatalError() point would also be reached if the function is called with an empty array. Actually that would also happen if the function is called with a non-empty array but the sum of the given weights is zero. These situations could be considered as a programming error on the caller's side and can be checked with a precondition().

@hashemi also suggested to use UInt instead of Int as data type for the weights. This is more natural because the weights must not be negative (and unexpected results would happen for negative input).

Then the function looks like this:

func weightedRandomElement<T>(items: [(T, UInt)]) -> T {

    let total = items.map { $0.1 }.reduce(0, combine: +)
    precondition(total > 0, "The sum of the weights must be positive")

    let rand = UInt(arc4random_uniform(UInt32(total)))

    var sum = UInt(0)
    for (element, weight) in items {
        sum += weight
        if rand < sum {
            return element
        }
    }

    fatalError("This should never be reached")
}

and is called like

let colors : [(String, UInt)] = [("Red", 20), ("Blue", 50), ("Green", 70)]
let color = weightedRandomElement(colors)