Edit: As @hashemi correctly noticed in a comment, this point is also
reached if the function is called with an empty array. Such invalid
input can be checked with a precondition
.
func weightedColor(input: [(String, Int)]) -> String {
precondition(input.count > 0, "weightedColor called with empty list")
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")
}
func weightedRandomElement<T>(items: [(T, Int)]) -> T {
precondition(items.count > 0, "weightedRandomElement called with empty list")
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")
}
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)