heap_map/
lib.rs

1#[cfg(test)]
2extern crate rand_08 as rand;
3mod tests;
4
5use malloc_size_of_derive::MallocSizeOf as DeriveMallocSizeOf;
6use std::{cmp::Ordering, collections::HashMap, fmt::Debug, hash};
7
8/// The `HeapMap` maintain a max heap along with a hash map to support
9/// additional `remove` and `update` operations.
10#[derive(DeriveMallocSizeOf)]
11pub struct HeapMap<K: hash::Hash + Eq + Copy + Debug, V: Eq + Ord + Clone> {
12    data: Vec<Node<K, V>>,
13    mapping: HashMap<K, usize>,
14}
15
16#[derive(Clone, DeriveMallocSizeOf)]
17pub struct Node<K, V: Eq + Ord> {
18    key: K,
19    value: V,
20}
21
22impl<K, V: Eq + Ord> Node<K, V> {
23    pub fn new(key: K, value: V) -> Self { Node { key, value } }
24}
25
26impl<K, V: Eq + Ord> PartialEq for Node<K, V> {
27    fn eq(&self, other: &Self) -> bool { self.value.eq(&other.value) }
28}
29
30impl<K, V: Eq + Ord> Eq for Node<K, V> {}
31
32impl<K, V: Eq + Ord> PartialOrd for Node<K, V> {
33    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
34        Some(self.cmp(other))
35    }
36}
37
38impl<K, V: Eq + Ord> Ord for Node<K, V> {
39    fn cmp(&self, other: &Self) -> Ordering { self.value.cmp(&other.value) }
40}
41
42impl<K: hash::Hash + Eq + Copy + Debug, V: Eq + Ord + Clone> Default
43    for HeapMap<K, V>
44{
45    fn default() -> Self { Self::new() }
46}
47
48impl<K: hash::Hash + Eq + Copy + Debug, V: Eq + Ord + Clone> HeapMap<K, V> {
49    pub fn new() -> Self {
50        Self {
51            data: vec![],
52            mapping: HashMap::new(),
53        }
54    }
55
56    /// Insert a K-V into the HeapMap.
57    /// Return the old value if `key` already exist. Return `None` otherwise.
58    pub fn insert(&mut self, key: &K, value: V) -> Option<V> {
59        if self.mapping.contains_key(key) {
60            let old_value = self.update(key, value);
61            Some(old_value)
62        } else {
63            self.append(key, value);
64            None
65        }
66    }
67
68    /// Remove `key` from the HeapMap.
69    pub fn remove(&mut self, key: &K) -> Option<V> {
70        let index = self.mapping.remove(key)?;
71        let removed_node = self.data.swap_remove(index);
72        if index != self.data.len() {
73            // The last node has been swapped to index
74            match self.data[index].cmp(&removed_node) {
75                Ordering::Less => self.sift_down(index),
76                Ordering::Greater => self.sift_up(index),
77                Ordering::Equal => {
78                    self.mapping.insert(self.data[index].key, index);
79                }
80            }
81        }
82        Some(removed_node.value)
83    }
84
85    /// In-place update some fields of a node's value.
86    pub fn update_with<F>(&mut self, key: &K, mut update_fn: F)
87    where F: FnMut(&mut V) {
88        let index = match self.mapping.get(key) {
89            None => {
90                return;
91            }
92            Some(i) => *i,
93        };
94        let origin_node = self.data[index].clone();
95        update_fn(&mut self.data[index].value);
96        // The order of node is the opposite of the order of this tuple.
97        match self.data[index].cmp(&origin_node) {
98            Ordering::Less => self.sift_down(index),
99            Ordering::Greater => self.sift_up(index),
100            _ => {}
101        }
102    }
103
104    /// Return the top K-V reference tuple.
105    pub fn top(&self) -> Option<(&K, &V)> {
106        self.data.first().map(|node| (&node.key, &node.value))
107    }
108
109    /// Pop the top node and return it as a K-V tuple.
110    pub fn pop(&mut self) -> Option<(K, V)> {
111        if self.is_empty() {
112            return None;
113        }
114        let item = self.data.swap_remove(0);
115        if !self.is_empty() {
116            self.sift_down(0);
117        }
118        self.mapping.remove(&item.key);
119        Some((item.key, item.value))
120    }
121
122    /// Get the value reference of `key`.
123    pub fn get(&self, key: &K) -> Option<&V> {
124        let index = *self.mapping.get(key)?;
125        self.data.get(index).map(|node| &node.value)
126    }
127
128    /// Clear all key-values of the HeapMap.
129    pub fn clear(&mut self) {
130        self.mapping.clear();
131        self.data.clear();
132    }
133
134    #[inline]
135    pub fn is_empty(&self) -> bool { self.data.is_empty() }
136
137    #[inline]
138    pub fn len(&self) -> usize { self.data.len() }
139
140    pub fn iter(&self) -> impl Iterator<Item = V> + '_ {
141        self.data.iter().map(|f| f.value.clone())
142    }
143
144    fn update(&mut self, key: &K, value: V) -> V {
145        let index = *self.mapping.get(key).unwrap();
146        let origin_node = self.data[index].clone();
147        self.data[index] = Node::new(*key, value);
148        match self.data[index].cmp(&origin_node) {
149            Ordering::Less => self.sift_down(index),
150            Ordering::Greater => self.sift_up(index),
151            _ => {}
152        }
153        origin_node.value
154    }
155
156    fn append(&mut self, key: &K, value: V) {
157        self.data.push(Node::new(*key, value));
158        self.sift_up(self.data.len() - 1);
159    }
160
161    fn sift_up(&mut self, index: usize) {
162        let val = self.data[index].clone();
163        let mut pos = index;
164        while pos > 0 {
165            let parent = (pos - 1) / 2;
166            if self.data[parent] >= val {
167                break;
168            }
169            self.data[pos] = self.data[parent].clone();
170            self.mapping.insert(self.data[pos].key, pos);
171            pos = parent;
172        }
173
174        self.mapping.insert(val.key, pos);
175        self.data[pos] = val;
176    }
177
178    fn sift_down(&mut self, index: usize) {
179        let val = self.data[index].clone();
180        let mut pos = index;
181        let mut child = pos * 2 + 1;
182        while child < self.data.len() {
183            let right = child + 1;
184            if right < self.data.len() && self.data[right] > self.data[child] {
185                child = right;
186            }
187            if val >= self.data[child] {
188                break;
189            }
190            self.data[pos] = self.data[child].clone();
191            self.mapping.insert(self.data[pos].key, pos);
192            pos = child;
193            child = pos * 2 + 1;
194        }
195        self.mapping.insert(val.key, pos);
196        self.data[pos] = val;
197    }
198}
heap_map/lib.rs

heap_map/
lib.rs
