按权重随机选择（leetcode 528）

文章目录

• 1.问题描述
• 2.难度等级
• 3.热门指数
• 4.解题思路
• 5.实现示例
  • 5.1 C++
  • 5.2 Golang
• 参考文献

1.问题描述

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

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

要求: 返回的随机函数的时间复杂度不超过 O(logn)。

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

2.难度等级

medium。

3.热门指数

★★★★☆

出题公司：萌时科技。

4.解题思路

可以使用“前缀和 + 二分查找”来实现。

设数组 w 的权重之和为 total。根据题目的要求，我们可以看成将 [1,total] 范围内的所有整数分成 n 个部分（其中 n 是数组 w 的长度）。第 i 个部分恰好包含 w[i] 个整数，并且这 n 个部分两两的交集为空。随后我们在 [1,total] 范围内随机生成一个整数 x，如果整数 x 被包含在第 i 个部分内，我们就返回 i。

例如对于数组 w[1,3,5,6]，计算其累计的前缀和数组为 [1,4,9,15]，然后随机产生一个 [1,15] 之间的随机数。

如果随机数落在 [1,1]，应该找到的值为 1, 对应元素下标为 0， 如果随机数落在 [2,4] 区间，应该找到值 4, 对应元素下标为 1， 如果随机数落在 [5,9] 区间，应该找到值 9, 对应元素下标为 2， 如果随机数落在 [10,15]，应该找到值 15, 对应元素下标为 3，

如果使用顺序遍历来查找元素效率较低, 由于前缀和数组是有序的, 所以可以使用二分法查找。

注意：需要先确定待查找的元素在数组中的位置 , 根据题意, 前缀和数组中第一个 >= x 的元素就是待查找的元素，确定了这一点才能写出正确的二分查找。

复杂度分析：

时间复杂度：初始化的时间复杂度为 O(n)，每次选择的时间复杂度为 O(logn)，其中 n 是数组 w 的长度。

空间复杂度：O(n)，即前缀和数组需要使用的空间。

5.实现示例

5.1 C++

#include <random>
#include <vector>
#include <numeric>
#include <iterator>
using namespace std;

class Solution {
private:
    mt19937 gen;
    uniform_int_distribution<int> dis;
    vector<int> pre;

public:
    Solution(vector<int>& w): gen(random_device{}()), dis(1, accumulate(w.begin(), w.end(), 0)) {
        partial_sum(w.begin(), w.end(), back_inserter(pre));
    }
    
    int pickIndex() {
        int x = dis(gen);
        return lower_bound(pre.begin(), pre.end(), x) - pre.begin();
    }
};

5.2 Golang

import (
	"sort"
	"math/rand"
)

type Solution struct {
    pre []int
}

func Constructor(w []int) Solution {
    for i := 1; i < len(w); i++ {
        w[i] += w[i-1]
    }
    return Solution{w}
}

func (s *Solution) PickIndex() int {
    x := rand.Intn(s.pre[len(s.pre)-1]) + 1
    return sort.SearchInts(s.pre, x)
}

参考文献

528. 按权重随机选择 - leetcode