1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

// Copyright 2019 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

use rand::{prelude::ThreadRng, Rng};
use std::{collections::HashMap, hash::Hash, slice::Iter};

/// HashMap that provide sampling in O(1) complexity.
#[derive(Default, Debug)]
pub struct SampleHashMap<K: Hash + Eq, V> {
    items: Vec<(K, V)>,
    index: HashMap<K, usize>,
}

impl<K: Hash + Eq + Clone, V> SampleHashMap<K, V> {
    pub fn get(&self, k: &K) -> Option<&V> {
        let pos = self.index.get(k)?;
        Some(&self.items[*pos].1)
    }

    pub fn get_mut(&mut self, k: &K) -> Option<&mut V> {
        let pos = self.index.get(k)?;
        Some(&mut self.items[*pos].1)
    }

    pub fn get_mut_or_insert_with<F: FnOnce() -> V>(
        &mut self, k: K, default: F,
    ) -> &mut V {
        let pos = match self.index.get(&k) {
            Some(pos) => *pos,
            None => {
                self.index.insert(k.clone(), self.items.len());
                self.items.push((k, default()));
                self.items.len() - 1
            }
        };

        &mut self.items[pos].1
    }

    pub fn remove(&mut self, k: &K) -> Option<V> {
        let index = self.index.remove(k)?;
        let (_, removed) = self.items.swap_remove(index);

        if let Some((swapped, _)) = self.items.get(index) {
            self.index.insert(swapped.clone(), index);
        }

        Some(removed)
    }

    pub fn sample(&self, rng: &mut ThreadRng) -> Option<&V> {
        if self.items.is_empty() {
            return None;
        }

        let index = rng.gen_range(0, self.items.len());
        Some(&self.items[index].1)
    }

    pub fn is_empty(&self) -> bool { self.items.is_empty() }
}

/// HashSet that provide sampling in O(1) complexity.
#[derive(Default, Debug)]
pub struct SampleHashSet<T: Hash + Eq> {
    items: Vec<T>,
    index: HashMap<T, usize>,
}

impl<T: Hash + Eq + Clone> SampleHashSet<T> {
    pub fn insert(&mut self, value: T) -> bool {
        if self.index.contains_key(&value) {
            return false;
        }

        self.index.insert(value.clone(), self.items.len());
        self.items.push(value);

        true
    }

    pub fn remove(&mut self, value: &T) -> bool {
        let index = match self.index.remove(value) {
            Some(pos) => pos,
            None => return false,
        };

        self.items.swap_remove(index);

        if let Some(swapped) = self.items.get(index) {
            self.index.insert(swapped.clone(), index);
        }

        true
    }

    pub fn sample(&self, rng: &mut ThreadRng) -> Option<T> {
        if self.items.is_empty() {
            return None;
        }

        let index = rng.gen_range(0, self.items.len());
        Some(self.items[index].clone())
    }

    #[inline]
    pub fn is_empty(&self) -> bool { self.items.is_empty() }

    #[inline]
    pub fn len(&self) -> usize { self.items.len() }

    #[inline]
    pub fn iter(&self) -> Iter<T> { self.items.iter() }
}