洗牌算法详解

14 阅读1分钟

前言

前段时间在准备求职面试时,在一个销售公司的面试中,多次被问及如何实现一个洗牌算法。我对洗牌算法的了解仅限于在 Radash 函数库中见过相关函数,对此并不熟悉,于是后面就看了一下相关的代码实现,来学习一下。

洗牌算法的核心目标是将一个序列的元素顺序完全随机化,确保每种可能的排列出现的概率相等。本文将对几种经典的洗牌算法进行介绍。

Fisher-Yates

Fisher-Yates 洗牌算法由 Ronald Fisher 和 Frank Yates 于1938年提出,是一种时间复杂度为 O(n)O(n) 的洗牌算法。其基本思想是从序列的最后一个元素开始,随机选择一个从第一个元素到当前位置之间的元素进行交换,然后逐步向前移动,直到整个序列被遍历一次。

代码实现

它也有前向和原地交换的版本,原理上都一样的。

function fisherYatesShuffle<T>(array: T[]): T[] {
  const shuffled = [...array]
  for (let i = shuffled.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * (i + 1))
    [shuffled[i], shuffled[j]] = [shuffled[j], shuffled[i]]
  }
  return shuffled
}

正确性证明

考虑一个长度为 nn 的序列。算法的核心是从后向前处理每个位置:对于位置 ii(从 n1n-1 开始递减到 1),我们从 0 到 ii 之间随机选择一个位置 jj,然后交换位置 ii 和位置 jj 的元素。

现在,我们关注任意一个特定元素最终出现在任意特定位置的概率。以元素 xx 最终出现在位置 kk 为例,看看这个概率是多少。

算法执行过程中,元素 xx 最终被放置到位置 kk,需要满足两个条件:

  • 在处理的每一步(从位置 n1n-1 到位置 k+1k+1),元素 xx 都没有被选中与当前处理的 ii 位置交换;
  • 在轮到处理位置 kk 时,元素 xx 被选中与位置 kk 交换。

让我们计算这两个条件同时满足的概率:

当处理最后一个位置 n1n-1 时,算法从所有 nn 个元素中随机选择一个放到这个位置。元素 xx 不被选中的概率是 (n1)/n(n-1)/n

接着处理位置 n2n-2,此时剩余 n1n-1 个元素(因为位置 n1n-1 已确定),算法从这 n1n-1 个元素中随机选择一个放到位置 n2n-2。元素 xx 不被选中的概率是 (n2)/(n1)(n-2)/(n-1)

依此类推,当处理到位置 k+1k+1 时,剩余 k+2k+2 个元素,元素 xx 不被选中的概率是 (k+1)/(k+2)(k+1)/(k+2)

最后,当处理到位置 kk时,剩余 k+1k+1 个元素,算法从这 k+1k+1 个元素中随机选择一个放到位置 kk。这时元素 xx 被选中的概率是 1/(k+1)1/(k+1)

将所有概率相乘:

(n1)/n×(n2)/(n1)×...×(k+1)/(k+2)×1/(k+1)=1/n(n-1)/n \times (n-2)/(n-1) \times ... \times (k+1)/(k+2) \times 1/(k+1) = 1/n

这意味着,任意一个元素最终出现在任意特定位置的概率都是 1/n1/n

此外,可以证明对于任意 n,n1n,n \geq 1,Fisher-Yates shuffle 生成任意特定排列 π\pi 的概率为 1/n!1/n!

n=1n=1,自然成立。假设对于 n1n-1 长度的数组,算法生成某个任意的排列的概率为 1/(n1)!1/(n−1)!

考虑长度为 nn 的数组:

  1. i=n1i=n-1,随机选择 j[0,n1]j \in [0,n-1],交换下标 jjn1n-1,这意味着位置 n1n-1 上的元素是从所有 nn 个元素中随即均匀选择的。对于任意特定元素,xx 选到 n1n-1 上的概率是 1/n1/n
  2. ii 继续向前移动,接下来的步骤是对 [0,n2][0,n-2] 上的元素继续使用 Fisher-Yates shuffle。根据归纳假设,这 n1n-1 个元素生成特定排列的概率为 1/(n1)!1/(n−1)!

上述过程不依赖特定的 xxπ\pi,因此,Fisher-Yates shuffle 生成任意特定排列 ππ 的概率为 1/n!1/n!

Sattolo 算法

Sattolo 算法是 Fisher-Yates 算法的一个变种,由 Sandra Sattolo 于1986年提出。它与 Fisher-Yates 算法的唯一区别在于:在每一步中,随机选择的索引 jj 的范围是 [0,i1][0, i-1],而不是 [0,i][0, i]。这个细微的修改导致算法产生的排列是一个随机的循环排列,即每个元素都不停留在其原始位置上。

代码实现

一样的,也有前向和原地修改的版本。

function sattoloShuffle<T>(array: T[]): T[] {
  const shuffled = [...array]
  for (let i = shuffled.length - 1; i > 0; i--) {
    const j = Math.floor(Math.random() * i)
    [shuffled[i], shuffled[j]] = [shuffled[j], shuffled[i]]
  }
  return shuffled
}

正确性证明

Sattolo 算法生成的是一个长度为 nn 的随机循环排列。所谓循环排列,是指排列可以被表示为若干个不相交的循环,且整个排列本身构成一个单一的循环(即从任一元素开始,通过重复应用排列映射,可以遍历所有元素后回到起点)。

类似地,我们可以知道,对于每一个下标 ii 的元素,放在某个下标 k,kik, k \not= i 中的概率是

(n2)/(n1)×(n3)/(n2)×...×k/(k+1)×1/k=1/(n1),k>0(n2)/(n1)×(n3)/(n2)×...×2/3×1/2=1/(n1),k=0(n-2)/(n-1) \times (n-3)/(n-2) \times ... \times k/(k+1) \times 1/k = 1/(n-1),k>0\\ (n-2)/(n-1) \times (n-3)/(n-2) \times ... \times 2/3 \times 1/2 = 1/(n-1),k=0

在算法的每一步中,由于下标 jj 的取值范围排除了 ii 本身,这确保了下标 ii 元素不会与自身交换,从而不会被放回原位。更重要的是,这种选择方式强制在排列图中形成一条边从 ii 指向 jj,最终这些边必然构成一个单一的循环。

可以证明,Sattolo 算法均匀地生成所有可能的 (n1)!(n-1)! 个循环排列。直观理解是,在第一步(处理索引 n1n-1),我们有 n1n-1 种选择来建立一条边;在后续的每一步,可用的选择数递减,但整体上每个循环排列的生成路径都是唯一且等概率的,概率为

1/(n1)×1/(n2)×...×1/1=1/(n1)!1/(n-1) \times 1/(n-2) \times ... \times 1/1 = 1/(n-1)!

蓄水池抽样算法

蓄水池抽样算法用于解决这样一个问题:当数据的总量 n 非常大或者未知时,如何从中等概率地随机抽取 kk 个样本。该算法只需对数据进行一次遍历,且仅需 O(k)O(k) 的额外内存空间。当 k=1k=1 时,可以看作是对整个“数据流”进行一种特殊的“洗牌”,最终随机获得一个元素;当 k>1k>1 时,它可以看作是从流动的数据中动态维护一个随机样本集。

代码实现

function reservoirSample<T>(stream: Iterable<T>, k: number): T[] {
  const reservoir: T[] = []
  let count = 0
  
  for (const item of stream) {
    count++;
    if (count <= k) {
      reservoir.push(item);
    } else {
      const replaceIndex = Math.floor(Math.random() * count)
      if (replaceIndex < k) {
        reservoir[replaceIndex] = item
      }
    }
  }
  return reservoir
}

正确性证明

容量为 1 时

我们首先证明 k=1k=1 的情况的正确性。算法维护一个变量 reservoirreservoir 作为当前的候选样本。当处理第 ii 个元素(ii 从 1 开始计数)时,算法以 1/i1/i 的概率选择用新元素替换 reservoirreservoir 中的当前值。

对于任意一个索引为 mm 的元素,它最终被选为样本的概率,等于它在第 mm 步被选入蓄水池的概率,乘以在后续所有步骤中都没有被替换出去的概率。具体计算为:

1m×mm+1×...×n1n=1n\frac{1}{m} \times \frac{m}{m+1} \times... \times\frac{n-1}n=\frac{1}n

因此,每个元素被选中的概率都是相等的 1/n1/n

容量大于 1 时

对于一般的 k>1k>1 的情况,证明思路类似。在处理第 ii 个元素时(i>ki > k),算法以 k/ik/i 的概率决定是否将其纳入蓄水池,如果决定纳入,则再均匀随机地替换掉蓄水池中现有的一个元素。通过类似的概率链式乘法,可以证明每个元素最终出现在蓄水池中的概率均为 k/nk/n

对于第 m,m>km,m>k 个元素,最终被选中的概率等于它在第 mm 步被选入蓄水池的概率,乘以在后续所有步骤中它没有被替换出去的概率。

在第 mm 步,它被选入的概率是 k/mk/m

一旦被选入,它会在蓄水池中占据一个位置。后续每处理一个新元素(第 j,j>mj,j>m 步),该新元素有 k/jk/j 的概率被选入,并且如果被选入,它会以 1/k1/k 的概率替换掉当前元素。因此,在第 jj 步,当前元素被替换的概率为 kj×1k=1j\frac{k}j \times\frac{1}k=\frac{1}j。那么它存活(不被替换)的概率就是 j1j\frac{j-1}j

所以,从第 m+1m+1 步到第 nn 步,它都存活下来的概率为:

mm+1×m+1m+2×...×n1n=mn\frac{m}{m+1} \times \frac{m+1}{m+2} \times... \times\frac{n-1}n=\frac{m}n

因此,第 mm 个元素最终被选中的概率为:mn×km=kn\frac{m}n\times\frac{k}m=\frac{k}{n}

对于前 kk 个元素(即 mkm≤k),这些元素一开始就在蓄水池中,但后续可能被替换。它们最终被选中的概率等于它们在后续所有步骤中从未被替换的概率。

对于第 m,mkm,m≤k 个元素,它从第 k+1k+1 步开始面临被替换的风险。同样,在第 j,j>kj,j>k 步,它被替换的概率为 1/j1/j ,存活概率为 j1j\frac{j-1}j

因此,从第 k+1k+1 步到第 nn 步,它一直存活下来的概率为:

kk+1×k+1k+2×...×n1n=kn\frac{k}{k+1}\times\frac{k+1}{k+2}\times...\times\frac{n-1}{n}=\frac{k}n

结语

洗牌算法虽然原理上并不复杂,但是它的设计极为巧妙。Fisher-Yates 算法提供了标准且高效的完全随机化方案,Sattolo 算法通过微小约束来满足完全错位的需求,而蓄水池抽样算法则可以突破数据必须全部在内存中的限制,优雅地处理流式数据。

大家的阅读是我发帖的动力,本文首发于我的博客:deer.shika-blog.xyz,欢迎大家来玩, 转载请注明出处。