加权随机：从离散概率分布中采样的算法

引言

首先，什么是加权随机？假设你有一个元素列表，想从中随机挑选一个。这看起来很简单，只需写一个小函数，生成一个指向列表中某个元素的随机下标即可。但有时单纯的随机并不够，我们想要带有偏向、或者依据某种概率的随机结果。这正是加权随机生成算法派上用场的地方。

应用场景

现实世界中有很多需要加权随机的场景，比如负载均衡器（例如 nginx、haproxy 等）。下面是一段 nginx 的示例配置。在这个例子中，backend1 的权重是 5；另外两台服务器采用默认权重 1。按这样的权重配置，每七个请求中，有五个发往 backend1，一个发往 backend2，一个发往 backend3。

http {
    upstream backend {
        server backend1 weight=5;
        server backend2;
        server backend3;
    }
}

另一个例子是爬虫调度。去年我在开发一个并发爬虫框架时，需要根据任务的重要程度来调度爬取任务。任务的重要程度用浮点数权重来表示，这些权重是手动为每个任务所属的站点分配的。因此，爬虫的调度器需要一个 WeightedChoice 函数，用来决定下一次该调度哪个任务。

在著名的 word2vec 的负采样部分，算法需要根据词在语料库中的频率来随机采样一些负样本词。代码链接。

游戏开发里还有更多例子：物品往往按设定的掉落概率随机产出，比如银币 25%、金币 20%、钻石 10%、装备 5%、饰品 40%，也就是说下一次掉落的物品类型要符合上述概率。

解决方案

方案一

我最先想到的方法，是扩展均匀分布的随机数生成器。我们先从一个例子开始（此后所有示例程序都用 Golang 编写）：

var items   = []int{0, 1, 2, 3}
var weights = []float32{0.1, 0.3, 0.4, 0.2}

要让它们符合这些权重，做法很简单：按权重的数值把每个元素重复 10 倍甚至 100 倍。假设重复 10 倍，最终得到的列表如下：

var choices = []int{0, 1, 1, 1, 2, 2, 2, 2, 3, 3}

然后我们就可以从 choices 中随机选一个值，大功告成。完整代码如下：

package main

import (
    "fmt"
    "math/rand"
)

func WeightedRandomS1(weights []float32) int {
    if len(weights) == 0 {
        return 0
    }
    var choices []int
    for i, w := range weights {
        wi := int(w * 10)
        for j := 0; j < wi; j++ {
            choices = append(choices, i)
        }
    }
    return choices[rand.Int()%len(choices)]
}

func main() {
    for i := 0; i < 100; i++ {
        fmt.Println(WeightedRandom([]float32{0.1, 0.3, 0.6}))
    }
}

在 word2vec 中，采用的正是这个方案。

方案二

第一个方案占用的内存太多，于是有了方案二：计算列表的离散累积分布函数（CDF）——用通俗的话说，就是权重的累加和数组。然后在 0 到所有权重之和这个区间内生成一个随机数，在离散 CDF 数组中做线性查找定位这个随机数，取出与该项对应的值——这就是加权随机数。

func WeightedRandomS2(weights []float32) int {
    if len(weights) == 0 {
        return 0
    }
    var sum float32 = 0.0
    for _, w := range weights {
        sum += w
    }
    r := rand.Float32() * sum
    for i, w := range weights {
        r -= w
        if r < 0 {
            return i
        }
    }
    return len(weights) - 1
}

我在爬虫框架的调度器里用的就是这个方案。

方案三

对 CDF 数组采用二分查找，可以把时间复杂度从 $O(n)$ 降到 $O(log(n))$ 。

func WeightedRandomS3(weights []float32) int {
    n := len(weights)
    if n == 0 {
        return 0
    }
    cdf := make([]float32, n)
    var sum float32 = 0.0
    for i, w := range weights {
        if i > 0 {
            cdf[i] = cdf[i-1] + w
        } else {
            cdf[i] = w
        }
        sum += w
    }
    r := rand.Float32() * sum
    var l, h int = 0, n - 1
    for l <= h {
        m := l + (h-l)/2
        if r <= cdf[m] {
            if m == 0 || (m > 0 && r > cdf[m-1]) {
                return m
            }
            h = m - 1
        } else {
            l = m + 1
        }
    }
    return -1
}

方案四

加权随机的最优解应当是 Alias Method（别名方法）。它需要 $O(n)$ 时间来初始化、 $O(1)$ 时间完成一次选取，以及 $O(n)$ 内存。一个 golang 版本的实现在这里。

算法：Vose 的 Alias Method

初始化：

创建数组 $Alias$ 和 $Prob$ ，各自大小为 $n$ 。
创建两个工作列表 $Small$ 和 $Large$ 。
将每个概率乘以 $n$ 。
对每个缩放后的概率 $p_i$ $p_{i}$ ：
1. 若 $p_i<1$ ，把 $i$ 加入 $Small$ 。
2. 否则 $p_i \geqslant 1$ ，把 $i$ 加入 $Large$ 。
当 $Small$ $S ma l l$ 和 $Large$ $L a r g e$ 都非空时：（ $Large$ $L a r g e$ 可能会先被清空）
1. 从 $Small$ 中移除第一个元素，记为 $l$ 。
2. 从 $Large$ 中移除第一个元素，记为 $g$ 。
3. 令 $Prob[l]=p_l$ 。
4. 令 $Alias[l]=g$ 。
5. 令 $p_g = p_g + p_l - 1$ 。（这是一种数值上更稳定的写法）
6. 若 $p_g<1$ ，把 $g$ 加入 $Small$ 。
7. 否则 $p_g \geqslant 1$ ，把 $g$ 加入 $Large$ 。
当 $Large$ $L a r g e$ 非空时：
1. 从 $Large$ 中移除第一个元素，记为 $g$ 。
2. 令 $Prob[g] = 1$ 。
当 $Small$ $S ma l l$ 非空时：这只有在数值不稳定时才可能发生。
1. 从 $Small$ 中移除第一个元素，记为 $l$ 。
2. 令 $Prob[l] = 1$ 。

生成：

掷一枚 n 面的均匀骰子，得到某一面 $i$ 。
抛一枚有偏硬币，它以概率 $Prob[i]$ 正面朝上。
若硬币出现”正面”，返回 $i$ 。
否则，返回 $Alias[i]$ 。

参考文献

Walker, A. J. 1977. “An efficient method for generating discrete random variable with general distributions.” ACM Transactions on Mathematical Software 3 253–256.
“Darts, Dice, and Coins: Sampling from a Discrete Distribution”. Keith Schwarz, December 29, 2011
Alias method. Wikipedia, April 5, 2018
Weighted random selection from array , stackoverflow, Dec 16 2010