Medium

题目描述

给你一个 下标从 0 开始 的正整数数组 w，其中 w[i] 代表第 i 个下标的权重。

请你实现一个函数 pickIndex()，它可以 随机地 从范围 [0, w.length - 1]（含 0 和 w.length - 1）内按权重选择并返回一个下标。选择下标 i 的 概率 为 w[i] / sum(w)。

例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25（即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。

示例 1：

输入：
["Solution","pickIndex"]
[[[1]],[]]
输出：
[null,0]

解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。

示例 2：

输入：
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
输出：
[null,1,1,1,1,0]

解释：
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。

由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。

提示：

  • 1 <= w.length <= 10^4
  • 1 <= w[i] <= 10^5
  • pickIndex 将被调用不超过 10^4 次

解题思路

这道题是经典的加权随机选择问题。核心思想是构建前缀和数组，然后使用二分查找来实现按权重的随机选择。

算法思路：

  1. 构建前缀和数组：在构造函数中，计算权重数组的前缀和。前缀和数组中每个位置表示从开始到当前位置的权重总和。

  2. 随机选择：在 pickIndex() 中，生成一个 [1, totalSum] 范围内的随机数，然后使用二分查找在前缀和数组中找到第一个大于等于这个随机数的位置。

为什么这样有效？

  • 权重越大的元素在前缀和数组中占据的"区间"越长
  • 随机数落在某个区间的概率正好等于该区间长度与总长度的比值
  • 这正好对应了权重比例

举例说明： 对于 w = [1, 3]，前缀和为 [1, 4]：

  • 随机数为 1 时，找到第一个 ≥1 的位置是 0
  • 随机数为 2,3,4 时，找到第一个 ≥2,3,4 的位置都是 1
  • 因此选择下标 0 的概率是 1/4，选择下标 1 的概率是 3/4

推荐解法：前缀和 + 二分查找，时间复杂度最优。

代码实现

class Solution {
private:
    vector<int> prefixSums;
    
public:
    Solution(vector<int>& w) {
        prefixSums.resize(w.size());
        prefixSums[0] = w[0];
        for (int i = 1; i < w.size(); i++) {
            prefixSums[i] = prefixSums[i - 1] + w[i];
        }
    }
    
    int pickIndex() {
        int totalSum = prefixSums.back();
        int target = rand() % totalSum + 1;
        
        int left = 0, right = prefixSums.size() - 1;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (prefixSums[mid] >= target) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
};

import random

class Solution:
    def __init__(self, w: List[int]):
        self.prefix_sums = []
        prefix_sum = 0
        for weight in w:
            prefix_sum += weight
            self.prefix_sums.append(prefix_sum)
    
    def pickIndex(self) -> int:
        target = random.randint(1, self.prefix_sums[-1])
        
        left, right = 0, len(self.prefix_sums) - 1
        while left < right:
            mid = (left + right) // 2
            if self.prefix_sums[mid] >= target:
                right = mid
            else:
                left = mid + 1
        return left

public class Solution {
    private int[] prefixSums;
    private Random random;
    
    public Solution(int[] w) {
        prefixSums = new int[w.Length];
        prefixSums[0] = w[0];
        for (int i = 1; i < w.Length; i++) {
            prefixSums[i] = prefixSums[i - 1] + w[i];
        }
        random = new Random();
    }
    
    public int PickIndex() {
        int totalSum = prefixSums[prefixSums.Length - 1];
        int target = random.Next(1, totalSum + 1);
        
        int left = 0, right = prefixSums.Length - 1;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (prefixSums[mid] >= target) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
}

var Solution = function(w) {
    this.prefixSums = [];
    let prefixSum = 0;
    for (let weight of w) {
        prefixSum += weight;
        this.prefixSums.push(prefixSum);
    }
};

Solution.prototype.pickIndex = function() {
    const target = Math.floor(Math.random() * this.prefixSums[this.prefixSums.length - 1]) + 1;
    
    let left = 0, right = this.prefixSums.length - 1;
    while (left < right) {
        const mid = Math.floor((left + right) / 2);
        if (this.prefixSums[mid] >= target) {
            right = mid;
        } else {
            left = mid + 1;
        }
    }
    return left;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

操作时间复杂度空间复杂度
构造函数O(n)O(n)
pickIndexO(log n)O(1)
  • 时间复杂度：构造函数需要 O(n) 时间计算前缀和；pickIndex 使用二分查找，时间复杂度为 O(log n)
  • 空间复杂度：O(n)，用于存储前缀和数组

相关题目