intset/
lib.rs

1#![crate_name = "intset"]
2// Copyright 2020 Rob King
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! This crate provides a collection of data structures for storing sets of integers.
17//! The different data structures are designed to make different operations efficient.
18//!
19//! All of the data structures in this crate support the following operations with the
20//! associated complexity:
21//!
22//! - **contains** - check if an integer is in the set in O(1) time
23//! - **iterate** - iterate over the members of the set in O(*n*) time, where *n* is the number of
24//!                 elements in the set
25//! - **len** - return the number of elements in the set in O(1) time
26//!
27//! Individual set data structures support additional operations, as documented below.
28//!
29//! All of the set data structures in this crate have a maximum capacity, specified as the
30//! largest integer that can be stored in the set plus one. Once a set is created, it does
31//! no further allocations.
32
33/// A `GrowSet` is a set of integers that supports efficient addition and clearing.
34///
35/// It supports the following additional operations with the associated complexity:
36///
37/// - **add** - add an integer to the set in O(1) time
38/// - **clear** - remove all members from the set in O(1) time
39/// - **pop** - remove and return a random member of the set in O(1) time
40///
41/// `GrowSet`s are useful for sets that need to be cleared frequently and rebuilt.
42///
43/// The implementation of `GrowSet` is based on "An Efficient Represtation of Sparse Sets"
44/// by Briggs and Torczon (1993).
45#[derive(Debug)]
46pub struct GrowSet {
47    n: usize,
48    sparse: Vec<usize>,
49    dense: Vec<usize>,
50}
51
52impl<'a> GrowSet {
53    /// Returns a GrowSet with the given capacity.
54    ///
55    /// # Arguments
56    ///
57    /// * `size` - the set can store all integers up to, but not including, this value.
58    ///
59    /// # Examples
60    ///
61    /// ```
62    /// // This set can store the integers 0, 1, 2, 3, and 4.
63    /// use intset::GrowSet;
64    /// let mut set = GrowSet::with_capacity(5);
65    /// set.add(3);
66    /// assert!(set.contains(3));
67    /// ```
68    pub fn with_capacity(size: usize) -> Self {
69        Self {
70            n: 0,
71            sparse: vec![0; size],
72            dense: vec![0; size],
73        }
74    }
75
76    /// Returns one greater than the largest integer that can be stored in this set.
77    /// That is, a set that can store integers from 0 to 5 will have a capacity of 6.
78    pub fn capacity(&self) -> usize {
79        self.sparse.capacity()
80    }
81
82    /// Return the number of integers in the set.
83    ///
84    /// # Examples
85    ///
86    /// ```
87    /// use intset::GrowSet;
88    /// let mut set = GrowSet::with_capacity(5);
89    /// set.add(1);
90    /// set.add(4);
91    /// assert_eq!(set.len(), 2);
92    /// ```
93    pub fn len(&self) -> usize {
94        self.n
95    }
96
97    /// Returns true if the specified value is in the set.
98    ///
99    /// # Arguments
100    /// * `value` - the element to test
101    ///
102    /// # Examples
103    ///
104    /// ```
105    /// use intset::GrowSet;
106    /// let mut set = GrowSet::with_capacity(5);
107    /// set.add(1);
108    /// assert!(set.contains(1));
109    /// assert!(!set.contains(3));
110    /// ```
111    pub fn contains(&self, value: usize) -> bool {
112        value < self.sparse.len()
113            && self.sparse[value] < self.n
114            && self.dense[self.sparse[value]] == value
115    }
116
117    /// Remove all items from the set.
118    pub fn clear(&mut self) {
119        self.n = 0;
120    }
121
122    /// Add an item to the set. The item must be less than the capacity of the set.
123    ///
124    /// # Arguments
125    ///
126    /// * `value` - the value to store
127    ///
128    /// # Examples
129    ///
130    /// ```
131    /// use intset::GrowSet;
132    /// let mut set = GrowSet::with_capacity(100);
133    /// set.add(42);
134    /// assert!(set.contains(42));
135    /// ```
136    pub fn add(&mut self, value: usize) {
137        if !self.contains(value) {
138            self.dense[self.n] = value;
139            self.sparse[value] = self.n;
140            self.n += 1;
141        }
142    }
143
144    /// Returns an iterator over the values in the set.
145    /// Uniqueness is guaranteed; ordering is not.
146    pub fn iter(&'a self) -> impl Iterator<Item = &usize> + 'a {
147        self.dense.iter().take(self.n)
148    }
149
150    /// Returns a random element of the set in constant
151    /// time, or None if the set is empty.
152    pub fn pop(&mut self) -> Option<usize> {
153        if self.n == 0 {
154            None
155        } else {
156            self.n -= 1;
157            Some(self.dense[self.n])
158        }
159    }
160}
161
162#[cfg(test)]
163mod grow_set_tests {
164    use super::*;
165
166    #[test]
167    fn test_add() {
168        let mut set = GrowSet::with_capacity(6);
169        set.add(1);
170        set.add(3);
171        set.add(4);
172        assert!(!set.contains(0));
173        assert!(set.contains(1));
174        assert!(!set.contains(2));
175        assert!(set.contains(3));
176        assert!(set.contains(4));
177        assert!(!set.contains(6));
178    }
179
180    #[test]
181    fn test_repeated_add() {
182        let mut set = GrowSet::with_capacity(6);
183        set.add(1);
184        set.add(1);
185        set.add(1);
186        set.add(1);
187        set.add(1);
188        set.add(1);
189        set.add(1);
190        set.add(1);
191        set.add(1);
192        set.add(1);
193        set.add(1);
194        set.add(1);
195        set.add(1);
196        set.add(1);
197        set.add(1);
198        assert!(set.contains(1));
199    }
200    
201    #[test]
202    fn test_clear() {
203        let mut set = GrowSet::with_capacity(6);
204        set.add(3);
205        set.add(4);
206        assert!(!set.contains(0));
207        assert!(!set.contains(1));
208        assert!(!set.contains(2));
209        assert!(set.contains(3));
210        assert!(set.contains(4));
211        assert!(!set.contains(6));
212
213        set.clear();
214        assert!(!set.contains(0));
215        assert!(!set.contains(1));
216        assert!(!set.contains(2));
217        assert!(!set.contains(3));
218        assert!(!set.contains(4));
219        assert!(!set.contains(6));
220    }
221
222    #[test]
223    fn test_iter() {
224        let mut set = GrowSet::with_capacity(6);
225        set.add(0);
226        set.add(2);
227        set.add(4);
228
229        for i in set.iter() {
230            match i {
231                0 | 2 | 4 => assert!(true),
232                _ => unreachable!(),
233            }
234        }
235
236        set.clear();
237
238        for _ in set.iter() {
239            unreachable!();
240        }
241    }
242
243    #[test]
244    fn test_pop() {
245        let mut set = GrowSet::with_capacity(6);
246        set.add(0);
247        set.add(2);
248        set.add(4);
249
250        assert_eq!(set.len(), 3);
251        assert_eq!(set.pop(), Some(4));
252        assert_eq!(set.len(), 2);
253        assert_eq!(set.pop(), Some(2));
254        assert_eq!(set.len(), 1);
255        assert_eq!(set.pop(), Some(0));
256        assert_eq!(set.len(), 0);
257        assert_eq!(set.pop(), None);
258        assert_eq!(set.len(), 0);
259
260        set.add(4);
261        assert_eq!(set.len(), 1);
262        assert_eq!(set.pop(), Some(4));
263        assert_eq!(set.len(), 0);
264        assert_eq!(set.pop(), None);
265        assert_eq!(set.len(), 0);
266    }
267}
268
269/// A `ShrinkSet` is a set of integers that supports efficient removal and refilling.
270///
271/// `ShrinkSet`s are automatically initialized such that they contain all integers
272/// up to, but not including, their capacity. For example, a `ShrinkSet` with a capacity
273/// of 5 will contain the integers 0, 1, 2, 3, and 4 upon initialization.
274///
275/// A `ShrinkSet` supports the following additional operations with the associated
276/// time complexity:
277///
278/// - **remove** - remove an integer from the set in O(1) time
279/// - **refill** - adds all removed elements back into the set in O(1) time
280/// - **pop** - remove and return a random member of the set in O(1) time
281///
282/// `ShrinkSet`s are useful for situations where we want to prune search spaces or
283/// work queues (for example) and then reset them to their initial state efficiently.
284///
285/// The algorithm used by `ShrinkSet` is believed to be novel, but if there is existing
286/// literature, please let me know so I can reference it here.
287#[derive(Debug)]
288pub struct ShrinkSet {
289    p: usize,
290    map: Vec<usize>,
291    values: Vec<usize>,
292}
293
294impl<'a> ShrinkSet {
295    /// Returns a `ShrinkSet` with the given capacity.
296    ///
297    /// # Arguments
298    ///
299    /// * `size` - at creation time, the set will contain all integers up to, but not including, this value
300    ///
301    /// # Examples
302    ///
303    /// ```
304    /// // This set stores all integers 0..=4.
305    /// use intset::ShrinkSet;
306    /// let mut set = ShrinkSet::new(5);
307    /// assert!(set.contains(3));
308    /// assert!(set.contains(4));
309    /// ```
310    pub fn new(size: usize) -> Self {
311        Self {
312            p: size,
313            map: (0..size).collect(),
314            values: (0..size).collect(),
315        }
316    }
317
318    /// Returns one greater than the largest integer that can be stored in this set.
319    /// That is, a set that can store integers from 0 to 5 will have a capacity of 6.
320    pub fn capacity(&self) -> usize {
321        self.map.capacity()
322    }
323
324    /// Return the number of integers in the set.
325    ///
326    /// # Examples
327    ///
328    /// ```
329    /// use intset::ShrinkSet;
330    /// let mut set = ShrinkSet::new(5);
331    /// set.remove(1);
332    /// set.remove(4);
333    /// assert_eq!(set.len(), 3);
334    /// ```
335    pub fn len(&self) -> usize {
336        self.p
337    }
338
339    /// Returns true if the specified value is in the set.
340    ///
341    /// # Arguments
342    /// * `value` - the element to test
343    ///
344    /// # Examples
345    ///
346    /// ```
347    /// use intset::ShrinkSet;
348    /// let mut set = ShrinkSet::new(5);
349    /// set.remove(1);
350    /// assert!(!set.contains(1));
351    /// assert!(set.contains(3));
352    /// ```
353    pub fn contains(&self, value: usize) -> bool {
354        value < self.map.len() && self.map[value] < self.p
355    }
356
357    /// Refills the set by adding all removed elements back.
358    ///
359    /// # Examples
360    ///
361    /// ```
362    /// use intset::ShrinkSet;
363    /// let mut set = ShrinkSet::new(5);
364    /// set.remove(0);
365    /// set.remove(2);
366    /// assert!(set.contains(1));
367    /// assert!(set.contains(3));
368    /// assert!(set.contains(4));
369    /// assert!(!set.contains(0));
370    /// assert!(!set.contains(2));
371    ///
372    /// set.refill();
373    /// assert!(set.contains(1));
374    /// assert!(set.contains(3));
375    /// assert!(set.contains(4));
376    /// assert!(set.contains(0));
377    /// assert!(set.contains(2));
378    /// ```
379    pub fn refill(&mut self) {
380        self.p = self.map.len();
381    }
382
383    /// Returns an iterator over the values in the set.
384    /// Uniqueness is guaranteed; ordering is not.
385    pub fn iter(&'a self) -> impl Iterator<Item = &usize> + 'a {
386        self.values.iter().take(self.p)
387    }
388
389    /// Remove an element from the set.
390    ///
391    /// # Arguments
392    /// - `item` - the item to remove
393    ///
394    /// # Examples
395    ///
396    /// ```
397    /// use intset::ShrinkSet;
398    /// let mut set = ShrinkSet::new(5);
399    /// set.remove(0);
400    /// set.remove(2);
401    /// assert!(set.contains(1));
402    /// assert!(set.contains(3));
403    /// assert!(!set.contains(0));
404    /// assert!(!set.contains(2));
405    /// ```
406    pub fn remove(&mut self, item: usize) -> usize {
407        if self.contains(item) {
408            let item_index = self.map[item];
409            let last_item = self.values[self.p - 1];
410            let last_item_index = self.map[last_item];
411
412            self.values[last_item_index] = item;
413            self.values[item_index] = last_item;
414            self.map[last_item] = item_index;
415            self.map[item] = last_item_index;
416            self.p -= 1;
417        }
418
419        item
420    }
421
422    /// Remove and return a random element from the set
423    /// in constant time, or None if the set is empty.
424    pub fn pop(&mut self) -> Option<usize> {
425        if self.p == 0 {
426            None
427        } else {
428            Some(self.remove(self.values[0]))
429        }
430    }
431}
432
433#[cfg(test)]
434mod shrink_set_tests {
435    use super::*;
436
437    #[test]
438    fn test_new() {
439        let set = ShrinkSet::new(5);
440        assert!(set.contains(0));
441        assert!(set.contains(1));
442        assert!(set.contains(2));
443        assert!(set.contains(3));
444        assert!(set.contains(4));
445    }
446
447    #[test]
448    fn test_remove() {
449        let mut set = ShrinkSet::new(6);
450        assert_eq!(set.remove(1), 1);
451        assert_eq!(set.remove(3), 3);
452        assert_eq!(set.remove(5), 5);
453
454        assert!(set.contains(0));
455        assert!(!set.contains(1));
456        assert!(set.contains(2));
457        assert!(!set.contains(3));
458        assert!(set.contains(4));
459        assert!(!set.contains(5));
460    }
461
462    #[test]
463    fn test_refill() {
464        let mut set = ShrinkSet::new(6);
465        set.remove(1);
466        set.remove(3);
467        set.remove(5);
468
469        assert!(set.contains(0));
470        assert!(!set.contains(1));
471        assert!(set.contains(2));
472        assert!(!set.contains(3));
473        assert!(set.contains(4));
474        assert!(!set.contains(5));
475
476        set.refill();
477
478        assert!(set.contains(0));
479        assert!(set.contains(1));
480        assert!(set.contains(2));
481        assert!(set.contains(3));
482        assert!(set.contains(4));
483        assert!(set.contains(5));
484    }
485
486    #[test]
487    fn test_iter() {
488        let mut set = ShrinkSet::new(6);
489        set.remove(1);
490        set.remove(3);
491        set.remove(5);
492
493        for i in set.iter() {
494            match i {
495                1 | 3 | 5 => unreachable!(),
496                _ => assert!(true),
497            }
498        }
499
500        set.refill();
501
502        assert!(set.contains(0));
503        assert!(set.contains(1));
504        assert!(set.contains(2));
505        assert!(set.contains(3));
506        assert!(set.contains(4));
507        assert!(set.contains(5));
508    }
509
510    #[test]
511    fn test_pop() {
512        let mut set = ShrinkSet::new(4);
513        set.remove(3);
514
515        let i = set.pop();
516        assert!(i == Some(0) || i == Some(1) || i == Some(2));
517
518        let j = set.pop();
519        assert!(j == Some(0) || j == Some(1) || j == Some(2));
520        assert!(j != i);
521
522        let k = set.pop();
523        assert!(k == Some(0) || k == Some(1) || k == Some(2));
524        assert!(k != i && k != j);
525
526        assert_eq!(set.pop(), None);
527
528        set.refill();
529
530        set.remove(3);
531
532        let i = set.pop();
533        assert!(i == Some(0) || i == Some(1) || i == Some(2));
534
535        let j = set.pop();
536        assert!(j == Some(0) || j == Some(1) || j == Some(2));
537        assert!(j != i);
538
539        let k = set.pop();
540        assert!(k == Some(0) || k == Some(1) || k == Some(2));
541        assert!(k != i && k != j);
542
543        assert_eq!(set.pop(), None);
544    }
545}