最近有个需求，需要从列表中随机选择一个项目，但各个项目有一定权重，类似无保底的抽卡效果。这时候就羡慕Python，可以直接用

from numpy.random import choice
draw = choice(list_of_candidates, number_of_items_to_pick,
              p=probability_distribution)

实现

既然没有，就只能自己写了。原理很简单，就是将原始列表根据权重映射到线段上的各个区间中，然后线段上随机取一个点。

    fun <T> getRandomItemWeighted(list: List<T>, weights: List<Int>): T {
        require(list.isNotEmpty())
        require(list.size == weights.size)

        val weightSum = weights.sum()
        if (weightSum <= 0) {
            return list.random() // shouldn't reach here
        }

        var weight = 0
        val value = random.nextInt(weights.sum())
        list.forEachIndexed { index, t ->
            weight += weights[index]
            if (weight > value) {
                return t
            }
        }
        return list.random() // shouldn't reach here
    }

验证正确性

准备一个列表，其权重为一单调数列。统计各项出现次数，然后做回归分析。若实现正确，各项出现次数之比应大致等于权重之比，即与列表项下标呈线性相关关系。

fun main() {
    val counter = IntArray(10) { 0 }
    val list = listOf(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
    val weights = list.reversed()
    repeat(1_0000_0000) {
        RandomUtil.getRandomItemWeighted(list, weights).let {
            counter[it]++
        }
    }
    counter.mapIndexed { index, i -> "$index, $i" }
        .joinToString("\n")
        .let { println(it) }
}

线性拟合

最后修改于 2023-04-23