Weighted Random Picker Swift

WeightedRandom.swift

class WeightedRandom<T> {
    var elements: [Element]
    
    init(elements: [Element] = []) {
        self.elements = elements
    }
    
    struct Element {
        let item: T
        let weight: Double
    }
    
    func randomElement() -> T? {
        guard elements.count > 0 else { return nil }
        
        // Create an array to store the cumulative sum of weights
        var cumulativeWeights: [Double] = []
        var totalWeight: Double = 0.0

        // Calculate the cumulative sum of weights
        for element in elements {
            totalWeight += element.weight
            cumulativeWeights.append(totalWeight)
        }
        
        // Generate a random number between 0 and the total weight
        let randomValue = Double.random(in: 0..<totalWeight)

        // Find the element that corresponds to the random value
        var selectedElement: T?
        for i in 0..<elements.count {
            if randomValue <= cumulativeWeights[i] {
                selectedElement = elements[i].item
                break
            }
        }
        return selectedElement
    }
}

Usage:

let rand = WeightedRandom<String>(elements: [
    .init(item: "A", weight: 0.1),
    .init(item: "B", weight: 0.0),
    .init(item: "C", weight: 0.2),
    .init(item: "D", weight: 0.3),
    .init(item: "E", weight: 0.1),
])

for _ in 0..<20 {
    print(rand.randomElement()!)
}