如何从含重复元素的数组高效抽取M个不重复随机元素？

需求 #

• 抽取的元素必须唯一
• 待抽样数组可能包含重复元素
• 待抽样数组无需预先排序

原实现代码 #

#include <random>
#include <vector>
#include <algorithm>
#include <iostream>


const std::vector<int> sample(const std::vector<int>& input, size_t n) {
    std::random_device rd;
    std::mt19937 engine(rd());
    std::uniform_int_distribution<int> dist(0, input.size() - 1);
    
    std::vector<int> result;
    result.reserve(n);

    size_t id;
    do {
        id = dist(engine);
        if (std::find(result.begin(), result.end(), input[id]) == result.end())
            result.push_back(input[id]);
    } while (result.size() < n);

    return result;
}


int main() {
    std::vector<int> input{0, 0, 1, 1, 2, 2, 3, 3, 4, 4};

    std::vector<int> result = sample(input, 3);

    for (const auto& item : result)
        std::cout << item << ' ';
    std::cout << std::endl;
}

问题 #

上述算法效率较低，同时需要补充逻辑：判断数组中唯一元素数量是否不小于M，若不足则抽取所有唯一元素。

优化方案 #

通过std::unordered_set先对原数组去重，将去重后的集合转为vector，再用std::shuffle打乱顺序，最后根据唯一元素数量调整返回结果的大小：

const std::vector<int> sample(const std::vector<int>& input, size_t M) {
    std::unordered_set<int> rem_dups(input.begin(), input.end());
    if (rem_dups.size() < M) M = rem_dups.size();
    
    std::vector<int> result(rem_dups.begin(), rem_dups.end());
    std::mt19937 g(std::random_device{}());
    std::shuffle(result.begin(), result.end(), g);

    result.resize(M);
    return result;
}

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