Y N
Y N

Reputation: 841

C++. Weighted std::shuffle

Is there a way to do nice and elegant weighted shuffling using standard library? There is std::discrete_distribution. What I want is something like this:

std::vector<T> data { N elements };
std::vector<int> weights { N weights };
std::shuffle(std::begin(data), std::end(data), something based on discrete distribution);

Upvotes: 8

Views: 1556

Answers (2)

Dan Stahlke
Dan Stahlke

Reputation: 1469

Check out Weighted Random Sampling (2005; Efraimidis, Spirakis). If you create a list of values -pow(rand(0,1), weights[i]) and sort this, you will get what you want.

Equivalently (and slightly faster), this list can be created using an exponential distribution.

std::vector<size_t> weighted_shuffle(std::vector<double> const &weights, std::mt19937 &rng)
{
    //auto uniform_dist = std::uniform_real_distribution<double>();
    auto exp_dist = std::exponential_distribution<double>();

    std::vector<std::pair<double, size_t>> index_pairs;
    index_pairs.reserve(weights.size());
    for (size_t i=0; i<weights.size(); ++i) {
        double const p = weights[i];
        // from Efraimidis, Spirakis
        //index_pairs.emplace_back(-std::pow(uniform_dist(rng), 1.0/p), i);
        // equivalent and a bit faster
        //index_pairs.emplace_back(-std::log(uniform_dist(rng))/p, i);
        // equivalent and fastest
        index_pairs.emplace_back(exp_dist(rng)/p, i);
    }
    std::sort(index_pairs.begin(), index_pairs.end());

    std::vector<size_t> indices;
    indices.reserve(weights.size());
    for (auto const &[w, i] : index_pairs)
        indices.emplace_back(i);
    return indices;
}

Upvotes: 2

Bob__
Bob__

Reputation: 12759

If OP intent is to shuffle a list r of items

such that, given a list of weights w, the element a[i] with weight w[i] should be the first element of the random shuffle r with probability w[i]/sum(w).

As stated in the page linked by Severin Pappadeux:

Weighted random shuffling is the same as weighted random sampling from a list a without replacement. That is, choose with probability w[i]/sum(w) element a[i] from a. Store this element in a list r. Then, remove element a[i] from a and w[i] from w, and select a new element of the modified list a, and so on until a is empty.

I'am not aware of such an algorithm in the Standard Library, but a simple implementation could be:

#include <random>
#include <algorithm>
#include <iterator>

template <class D, class W, class URBG>
void weighted_shuffle
    ( D first, D last
    , W first_weight, W last_weight
    , URBG&& g )
{
    while (first != last and first_weight != last_weight)
    {
        std::discrete_distribution dd(first_weight, last_weight);
        auto i = dd(g);
        if ( i )
        {
            std::iter_swap(first, std::next(first, i));
            std::iter_swap(first_weight, std::next(first_weight, i));
        }
        ++first;
        ++first_weight;
    }
}

Live example HERE.

Upvotes: 8

Related Questions