patternswiftMinor

Weighted Probability Problem in Swift

Submitted by: @import:stackexchange-codereview·Mar 10, 2026·

Viewed 0 times

swiftweightedprobabilityproblem

Problem

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?

Solution

There is one error in your implementation: The test

if rand 

- 
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(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(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)

Code Snippets

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

fatalError("This should never be reached")

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")
}

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")
}

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

Context

StackExchange Code Review Q#112605, answer score: 3

Revisions (0)

No revisions yet.