动态概率列表中匹配随机数的O(1)索引查找算法咨询

输入数据 #

[
  {index: 0, probability: 0.20},
  {index: 1, probability: 0.10},
  {index: 2, probability: 0.40},
  {index: 3, probability: 0.25},
  {index: 4, probability: 0.05},
]

问题描述 #

生成一个[0,1)区间的随机数后，是否存在**O(1)**时间复杂度的算法，找到该随机数匹配的索引？

目前使用的是O(N)复杂度的实现，代码如下：

let cumulative = 0;
const r = Math.random();
for(const v of list){
    cumulative += v.probability;
    if(r < cumulative){
      return v.index;
    }
}

已确认该O(N)算法能满足需求，现询问是否有更高效的实现方案。

回答 #

不存在无需预处理的O(1)算法，但可以通过预处理阶段将查询操作优化到O(1)时间，其中最经典的方案是别名方法（Alias Method）。

别名方法的核心逻辑 #

1. 预处理阶段（时间复杂度O(N log N)）：

   • 将所有概率值放大N倍（N为元素总数），得到各元素的权重值。
   • 维护三个数组：存储每个槽主元素索引的数组、存储对应别名元素索引的数组、记录主元素在槽中占比的概率数组。
   • 通过贪心策略，将权重≥1的元素与权重<1的元素配对，填充到各个槽位中，完成预处理。

2. 查询阶段（O(1)时间）：

   • 生成两个随机数：第一个随机数选择目标槽位（范围0到N-1）；第二个随机数判断选择该槽的主元素还是别名元素（若随机数小于主元素占比则选主元素，否则选别名元素）。

更易实现的次优方案 #

如果觉得别名方法实现复杂，前缀和+二分查找是更简单的高效替代方案：

• 预处理阶段：计算概率的前缀和数组，时间复杂度O(N)。
• 查询阶段：用二分查找在前缀和数组中定位随机数对应的索引，时间复杂度O(log N)，比原O(N)方案效率更高。

示例代码如下：

// 预处理：生成前缀和数组
const prefixSums = [];
let sum = 0;
for (const v of list) {
  sum += v.probability;
  prefixSums.push({ index: v.index, sum });
}

// 查询阶段
const r = Math.random();
let left = 0, right = prefixSums.length - 1;
while (left < right) {
  const mid = Math.floor((left + right) / 2);
  if (prefixSums[mid].sum < r) {
    left = mid + 1;
  } else {
    right = mid;
  }
}
return prefixSums[left].index;

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