btree_multiset - [no description]


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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

#![feature(allocator_api, extend_one, btreemap_alloc)]
//! [`BTreeSet`](std::collections::BTreeSet) but with duplicates.
use std::alloc::{Allocator, Global};
use std::collections::BTreeMap;
use std::collections::btree_map::{self, Entry};
use std::fmt::Debug;
use std::iter::FlatMap;

use smolvec::SmolVec;

#[derive(Clone)]
/// Like a btree set, but with duplicates.
/// Currently immutable cause i havent thought about it yet.
pub struct BTreeMultiset<K, T, F, A: Allocator + Clone> {
    inner: BTreeMap<K, SmolVec<T>, A>,
    f: F,
}
impl<K: Ord, T: Debug, F: Fn(&T) -> K, A: Allocator + Clone> Debug
    for BTreeMultiset<K, T, F, A>
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_set().entries(self.iter()).finish()
    }
}
impl<K, T, F: Fn(&T) -> K, A: Allocator + Clone> IntoIterator
    for BTreeMultiset<K, T, F, A>
{
    type Item = T;

    type IntoIter = FlatMap<
        btree_map::IntoIter<K, SmolVec<T>, A>,
        SmolVec<T>,
        fn((K, SmolVec<T>)) -> SmolVec<T>,
    >;

    fn into_iter(self) -> Self::IntoIter {
        self.inner.into_iter().flat_map(|x| x.1)
    }
}
impl<K: Ord, T, F: Fn(&T) -> K, A: Allocator + Clone> Extend<T>
    for BTreeMultiset<K, T, F, A>
{
    #[inline]
    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
        iter.into_iter().for_each(move |x| {
            self.insert(x);
        });
    }

    #[inline]
    fn extend_one(&mut self, x: T) {
        self.insert(x);
    }
}
impl<K: Ord, T, F: Fn(&T) -> K> BTreeMultiset<K, T, F, Global> {
    pub fn new(f: F) -> Self {
        Self { inner: BTreeMap::new_in(Global), f }
    }
}
type Iter<'a, K, T> = FlatMap<
    btree_map::Iter<'a, K, SmolVec<T>>,
    &'a SmolVec<T>,
    fn((&'a K, &'a SmolVec<T>)) -> &'a SmolVec<T>,
>;
type Range<'a, K, T> = FlatMap<
    btree_map::Range<'a, K, SmolVec<T>>,
    &'a SmolVec<T>,
    fn((&'a K, &'a SmolVec<T>)) -> &'a SmolVec<T>,
>;
impl<K, T, F, A: Allocator + Clone> BTreeMultiset<K, T, F, A> {
    pub fn into_inner(self) -> BTreeMap<K, SmolVec<T>, A> {
        self.inner
    }
}
impl<K: Ord, T, F: Fn(&T) -> K, A: Allocator + Clone>
    BTreeMultiset<K, T, F, A>
{
    pub fn new_in(f: F, a: A) -> Self {
        Self { inner: BTreeMap::new_in(a), f }
    }
    /// Gets an iterator that visits the elements in the [`BTreeMultiset`] in ascending order.
    /// Similar to `range(..)`.
    pub fn iter<'a>(&'a self) -> Iter<'a, K, T> {
        self.inner.iter().flat_map(|x| x.1)
    }
    /// Adds a value to the multi set.
    pub fn insert(&mut self, t: T) -> bool {
        match self.inner.entry((self.f)(&t)) {
            Entry::Vacant(vacant_entry) => {
                vacant_entry.insert(SmolVec::from_iter([t]));
                false
            }
            Entry::Occupied(occupied_entry) => {
                occupied_entry.into_mut().push(t);
                true
            }
        }
    }

    /// Constructs a double-ended iterator over a sub-range of elements in the multiset.
    /// The simplest way is to use the range syntax `min..max`, thus `range(min..max)` will
    /// yield elements from min (inclusive) to max (exclusive).
    /// The range may also be entered as `(Bound<T>, Bound<T>)`, so for example
    /// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
    /// range from 4 to 10.
    ///
    /// # Panics
    ///
    /// Panics if range `start > end`.
    /// Panics if range `start == end` and both bounds are `Excluded`.
    pub fn range<'a>(
        &'a self,
        r: impl std::ops::RangeBounds<K>,
    ) -> Range<'a, K, T> {
        self.inner.range(r).flat_map(|x| x.1)
    }

    // pub fn range_mut<'a, R>(
    //     &'a mut self,
    //     r: impl std::ops::RangeBounds<K>,
    // ) -> FlatMap<
    //     RangeMut<'a, K, SmolVec<T>>,
    //     &'a mut SmolVec<T>,
    //     fn((&'a K, &'a mut SmolVec<T>)) -> &'a mut SmolVec<T>,
    // > {
    //     self.inner.range_mut(r).flat_map(|x| x.1)
    // }
}
#[test]
fn works() {
    let mut x = BTreeMultiset::new(|x: &u32| *x);

    x.insert(4);
    x.insert(3);
    x.insert(7);
    x.insert(3);
    assert_eq!(x.iter().copied().collect::<Vec<_>>(), vec![3, 3, 4, 7]);
    assert_eq!(x.range(0..=4).copied().collect::<Vec<_>>(), vec![3, 3, 4]);
}