我可以有效地从HashSet中随机抽样吗？

Question

我有一个std::collections::HashSet,我想要采样并删除一个均匀的随机元素.

目前,我正在做的是使用随机抽样索引rand.gen_range,然后迭代HashSet到该索引以获取元素.然后我删除所选元素.这有效,但效率不高.有没有一种有效的方法来随机抽样元素？

这是我的代码的简化版本:

use std::collections::HashSet;

extern crate rand;
use rand::thread_rng;
use rand::Rng;

let mut hash_set = HashSet::new();

// ... Fill up hash_set ...

let index = thread_rng().gen_range(0, hash_set.len());
let element = hash_set.iter().nth(index).unwrap().clone();
hash_set.remove(&element);

// ... Use element ...

Answer 1

唯一允许在恒定时间内进行均匀采样的数据结构是具有恒定时间索引访问的数据结构。HashSet不提供索引，因此您可以\xe2\x80\x99t 在恒定时间内生成随机样本。

\n

我建议首先将哈希集转换为哈希集Vec，然后从向量中采样。要删除一个元素，只需将最后一个元素移动到它的位置 \xe2\x80\x93 无论如何，向量中元素的顺序并不重要。

\n

如果您想以随机顺序使用集合中的所有元素，您还可以对向量进行一次洗牌，然后对其进行迭代。

\n

以下是在恒定时间内从 a 中删除随机元素的示例实现Vec：

\n

use rand::{thread_rng, Rng};\n\npub trait RemoveRandom {\n    type Item;\n\n    fn remove_random<R: Rng>(&mut self, rng: &mut R) -> Option<Self::Item>;\n}\n\nimpl<T> RemoveRandom for Vec<T> {\n    type Item = T;\n\n    fn remove_random<R: Rng>(&mut self, rng: &mut R) -> Option<Self::Item> {\n        if self.len() == 0 {\n            None\n        } else {\n            let index = rng.gen_range(0..self.len());\n            Some(self.swap_remove(index))\n        }\n    }\n}\n

\n

（游乐场）

\n

Answer 2

考虑 Sven Marnach 的答案，我想使用向量，但我还需要带有重复的恒定时间插入。然后我意识到我可以同时维护一个向量和一个集合，并确保它们始终具有相同的元素。这将允许恒定时间插入重复数据删除和恒定时间随机删除。

这是我最终得到的实现：

struct VecSet<T> {
    set: HashSet<T>,
    vec: Vec<T>,
}

impl<T> VecSet<T>
where
    T: Clone + Eq + std::hash::Hash,
{
    fn new() -> Self {
        Self {
            set: HashSet::new(),
            vec: Vec::new(),
        }
    }
    fn insert(&mut self, elem: T) {
        assert_eq!(self.set.len(), self.vec.len());
        let was_new = self.set.insert(elem.clone());
        if was_new {
            self.vec.push(elem);
        }
    }
    fn remove_random(&mut self) -> T {
        assert_eq!(self.set.len(), self.vec.len());
        let index = thread_rng().gen_range(0, self.vec.len());
        let elem = self.vec.swap_remove(index);
        let was_present = self.set.remove(&elem);
        assert!(was_present);
        elem
    }
    fn is_empty(&self) -> bool {
        assert_eq!(self.set.len(), self.vec.len());
        self.vec.is_empty()
    }
}

Answer 3

2023 年，实现 Rust 标准库的 HashSet 的hashbrown箱引入了 RawTable API，该 API 允许对 HashSet 的内部状态进行不安全的、较低级别的访问。使用这个API，我们可以直接实现从HashSet中随机删除：

use hashbrown::HashSet;
use rand::prelude::*;
use std::hash::Hash;

fn remove_random<T, R>(set: &mut HashSet<T>, rng: &mut R) -> Option<T>
where
    R: Rng,
    T: Eq + PartialEq + Hash,
{
    if set.is_empty() {
        return None;
    }
    // If load factor is under 25%, shrink to fit.
    // We need a high load factor to ensure that the sampling succeeds in a reasonable time,
    // and the table doesn't rebalance on removals.
    // Insertions can only cause the load factor to reach as low as 50%,
    // so it's safe to shrink at 25%.
    if set.capacity() >= 8 && set.len() < set.capacity() / 4 {
        set.shrink_to_fit();
    }
    let raw_table = set.raw_table_mut();
    let num_buckets = raw_table.buckets();
    // Perform rejection sampling: Pick a random bucket, check if it's full,
    // repeat until a full bucket is found.
    loop {
        let bucket_index = rng.gen_range(0..num_buckets);
        // Safety: bucket_index is less than the number of buckets.
        // Note that we return the first time we modify the table,
        // so raw_table.buckets() never changes.
        // Also, the table has been allocated, because set is a HashSet.
        unsafe {
            if raw_table.is_bucket_full(bucket_index) {
                let bucket = raw_table.bucket(bucket_index);
                let ((element, ()), _insert_slot) = raw_table.remove(bucket);
                return Some(element);
            }
        }
    }
}

该程序需要启用 hashbrown 箱的“原始”功能，如下所示Cargo.toml：

[dependencies]
hashbrown = { version = "^0.14.2", features = ["raw"] }
rand = "^0.8.5"