# Weighted Probability Problem in Swift

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)

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?

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)

• Two suggestions: (1) The fatalError will be reached if an empty array is passed. You can either guard against that at the top or change the fatalError message to give a clue to the consumer of function when they try to debug the crash. (2) Consider changing the tuple to (String, UInt) as the functions logic assumes that the weights passed are not negative. – hashemi Dec 4 '15 at 22:43
• @hashemi: Good points. Re (1): This can be checked with a precondition, I have added that. Re (2): Yes, I had thought about that as well. The disadvantage is that then the function cannot be called with simple integer literals anymore: colors =[("Red", UInt(20)), ...]. – Thanks for the feedback! – Martin R Dec 4 '15 at 23:13
• @MartinR you can specify the tuple's type explicitly to avoid the need to cast each number individually: let colors: [(String, UInt)] = [...]. Also, in real life situations, the numbers will likely be coming from user input or a database, so being careful with the types is even more important and literal input is less important. – hashemi Dec 4 '15 at 23:24
• @hashemi: You are right, and I have updated the answer according to your feedback. – Martin R Dec 5 '15 at 10:24
• FWIW, in The Swift Programming Language, Apple explicitly advises that you use Int, “even when the values to be stored are known to be non-negative.” They advise that one “[use] UInt only when you specifically need an unsigned integer type with the same size as the platform’s native word size.” – Rob Dec 28 '15 at 1:25