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}