C++ 加权随机抽样

发布时间 2023-03-22 21:12:53作者: 1bite


在做五子棋 AI 的时候,遇到要根据权重从一组数里边随机选出一个这个问题。这个问题恰好也是 leetcode 第 528 道题,使用 C++,标准库就有现成的:

#include <random>

int random_pick(double const* w, int k) {
    std::random_device rd;
    return std::discrete_distribution<int>{w, w+k}(rd);
}

若要自己实现,大致如下:

#include <numeric>
#include <random>
#include <algorithm>

int random_pick(double const* w, int k) {
    double *sum = (double*)alloca(k*sizeof(double)); //stack allocation
    std::partial_sum(w, w+k, sum);
    std::random_device rd;
    std::uniform_real_distribution<double> rnd(0, sum[k-1]);
    return std::upper_bound(sum, sum+k, rnd(rd)) - sum;
}