Swift中按指定概率/权重高效随机选择元素的最优方案

如果需要频繁执行按权重随机选择的操作，前缀和数组+二分查找是最优且最高效的实现方式——预处理阶段仅需O(n)时间，每次查询仅需O(logn)时间，完全适配高频调用场景。

具体实现步骤 #

1. 预处理前缀和数组：把每个元素的权重累加，生成前缀和数组。比如元素权重为[2, 3, 5]，对应的前缀和数组就是[2, 5, 10]，总权重为10。
2. 生成随机数：生成一个范围在0到总权重之间的随机浮点数（比如0到10之间）。
3. 二分查找定位元素：在前缀和数组中，找到第一个大于该随机数的元素索引，这个索引对应的原集合元素就是按权重选中的结果。

Swift代码示例 #

可以把逻辑封装成通用扩展，方便复用：

extension Collection {
    func randomElement<T: BinaryFloatingPoint>(byWeights weights: [T]) -> Element? {
        guard !isEmpty, weights.count == count else { return nil }
        
        // 计算前缀和数组，统一转为Double处理
        var prefixSums = [Double]()
        var currentSum = 0.0
        for weight in weights {
            currentSum += Double(weight)
            prefixSums.append(currentSum)
        }
        guard let totalWeight = prefixSums.last, totalWeight > 0 else { return nil }
        
        // 生成0到totalWeight之间的随机数
        let random = Double.random(in: 0..<totalWeight)
        
        // 二分查找定位目标索引
        var low = 0, high = prefixSums.count
        while low < high {
            let mid = (low + high) / 2
            if prefixSums[mid] <= random {
                low = mid + 1
            } else {
                high = mid
            }
        }
        return self[indices[low]]
    }
}

性能优化补充 #

如果权重是固定不变的，可以提前计算并缓存前缀和数组，避免每次调用重复计算，进一步提升高频场景下的执行效率。

内容的提问来源于stack exchange，提问作者Jay