Return to Revisions

There is one error in your implantation: 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 implementation 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>(list: [(T, Int)]) -> T {
    
    let total = UInt32(list.map { $0.1 }.reduce(0, combine: +))
    let rand = Int(arc4random_uniform(total))
    
    var sum = 0
    for (element, weight) in list {
        sum += weight
        if rand < sum {
            return element
        }
    }
    
    fatalError("This should never be reached")
}