competitive_library/structure/
binary_trie.rs

1//! BinaryTrie
2
3#[derive(Debug, Default, Clone, PartialEq, Eq)]
4struct Node {
5    children: [Option<usize>; 2],
6    count: u64,
7}
8impl Node {
9    #[inline]
10    fn new() -> Self {
11        Self {
12            children: [None; 2],
13            count: 0,
14        }
15    }
16}
17
18#[derive(Debug, Default, Clone, PartialEq, Eq)]
19pub struct BinaryTrie {
20    nodes: Vec<Node>,
21}
22impl BinaryTrie {
23    /// 構築
24    #[inline]
25    pub fn new() -> Self {
26        Self {
27            nodes: vec![Node::new()],
28        }
29    }
30
31    /// 値の挿入
32    #[inline]
33    pub fn insert(&mut self, x: u32) -> u64 {
34        self.insert_n(x, 1)
35    }
36    #[inline]
37    pub fn insert_n(&mut self, x: u32, n: u64) -> u64 {
38        if n == 0 {
39            return 0;
40        }
41        let mut node_index = 0;
42        for i in (0..32).rev() {
43            self.nodes[node_index].count += n;
44
45            node_index = match self.nodes[node_index].children[(x >> i & 1) as usize] {
46                Some(i) => i,
47                None => {
48                    self.nodes.push(Node::new());
49                    self.nodes[node_index].children[(x >> i & 1) as usize] =
50                        Some(self.nodes.len() - 1);
51                    self.nodes.len() - 1
52                }
53            };
54        }
55        self.nodes[node_index].count += n;
56        self.nodes[node_index].count
57    }
58
59    /// xのカウント
60    #[inline]
61    pub fn count(&self, x: u32) -> u64 {
62        let mut node_index = Some(0);
63
64        for i in (0..32).rev() {
65            if node_index.is_none() {
66                return 0;
67            }
68            node_index = self.nodes[node_index.unwrap()].children[(x >> i & 1) as usize];
69        }
70        if node_index.is_none() {
71            return 0;
72        }
73        self.nodes[node_index.unwrap()].count
74    }
75
76    /// x 未満の値のカウント
77    #[inline]
78    pub fn count_less(&self, x: u32) -> u64 {
79        self.inner_count_than(x, 1)
80    }
81
82    /// x を超える値のカウント
83    #[inline]
84    pub fn count_more(&self, x: u32) -> u64 {
85        self.inner_count_than(x, 0)
86    }
87    #[inline]
88    fn inner_count_than(&self, x: u32, bit: u32) -> u64 {
89        let mut node_index = Some(0);
90
91        let mut count = 0;
92        for i in (0..32).rev() {
93            if node_index.is_none() {
94                break;
95            }
96            if (x >> i & 1) == bit {
97                count += match self.nodes[node_index.unwrap()].children[(bit ^ 1) as usize] {
98                    Some(i) => self.nodes[i].count,
99                    None => 0,
100                }
101            }
102            node_index = self.nodes[node_index.unwrap()].children[(x >> i & 1) as usize];
103        }
104        count
105    }
106
107    /// 値の削除
108    #[inline]
109    pub fn erase(&mut self, x: u32) -> Option<()> {
110        if 1 > self.count(x) {
111            return None;
112        }
113        self.inner_erase(x, 1)
114    }
115
116    /// 値をすべて削除
117    #[inline]
118    pub fn erase_all(&mut self, x: u32) -> Option<()> {
119        let erase_count = self.count(x);
120        if erase_count == 0 {
121            return None;
122        }
123        self.inner_erase(x, erase_count)
124    }
125
126    /// 値を削除
127    /// 内部関数
128    #[inline]
129    fn inner_erase(&mut self, x: u32, erase_count: u64) -> Option<()> {
130        let mut node_index = Some(0);
131        for i in (0..32).rev() {
132            self.nodes[node_index?].count -= erase_count;
133            node_index = self.nodes[node_index?].children[(x >> i & 1) as usize];
134        }
135        self.nodes[node_index?].count -= erase_count;
136
137        Some(())
138    }
139
140    /// xor 後の最小値を求める
141    #[inline]
142    pub fn xor_min(&self, x: u32) -> Option<u32> {
143        let mut ans = 0;
144
145        let mut node_index = Some(0);
146        for i in (0..32).rev() {
147            let bit = {
148                let mut buff = (x >> i & 1) as usize;
149                if self.nodes[node_index.unwrap()].children[buff]
150                    .filter(|&index| self.nodes[index].count > 0)
151                    .is_none()
152                {
153                    buff ^= 1;
154                }
155                buff
156            };
157            ans ^= (bit as u32) << i;
158            node_index = self.nodes[node_index.unwrap()].children[bit];
159        }
160        Some(ans ^ x)
161    }
162
163    /// 最小値を求める
164    #[inline]
165    pub fn min(&self) -> Option<u32> {
166        self.xth_element(1)
167    }
168
169    /// 最大値を求める
170    #[inline]
171    pub fn max(&self) -> Option<u32> {
172        let max = self.size();
173        self.xth_element(max)
174    }
175    #[inline]
176    pub fn size(&self) -> u64 {
177        self.nodes[0].count
178    }
179    #[inline]
180    pub fn xth_element(&self, xth: u64) -> Option<u32> {
181        if self.size() < xth || xth == 0 {
182            return None;
183        }
184        let mut x = xth;
185        let mut ans = 0;
186        let mut node_index = Some(0);
187
188        for i in (0..32).rev() {
189            let count = if let Some(i) = self.nodes[node_index.unwrap()].children[0] {
190                self.nodes[i].count
191            } else {
192                0
193            };
194
195            let bit = if count >= x {
196                0
197            } else {
198                x -= count;
199                1
200            };
201            ans ^= (bit as u32) << i;
202            node_index = self.nodes[node_index.unwrap()].children[bit];
203        }
204
205        Some(ans)
206    }
207
208    #[inline]
209    pub fn lower_bound(&self, x: u32) -> Option<u32> {
210        self.xth_element(self.count_less(x + 1))
211    }
212
213    #[inline]
214    pub fn upper_bound(&self, x: u32) -> Option<u32> {
215        self.xth_element(self.count_less(x + 1) + 1)
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn bt() {
225        let mut b = BinaryTrie::new();
226        b.insert(6);
227        assert_eq!(b.size(), 1);
228
229        b.insert(7);
230        b.insert(7);
231        assert_eq!(b.size(), 3);
232        assert_eq!(b.xth_element(1).unwrap(), 6);
233        assert_eq!(b.xth_element(2).unwrap(), 7);
234        assert_eq!(b.xth_element(3).unwrap(), 7);
235        b.erase(7);
236        b.erase(7);
237        assert_eq!(b.count(2), 0);
238        assert_eq!(b.count(3), 0);
239        assert_eq!(b.count(4), 0);
240        assert_eq!(b.count(5), 0);
241        assert_eq!(b.count(8), 0);
242        assert_eq!(b.size(), 1);
243        assert_eq!(b.erase(10), None);
244    }
245    #[test]
246    fn btt() {
247        let mut b = BinaryTrie::new();
248        let n = 2u32.pow(30);
249        for i in 0..100 {
250            b.insert(n + i);
251        }
252        for i in 0..99 {
253            b.erase(n + i);
254            assert_eq!(b.min().unwrap(), n + i + 1);
255        }
256    }
257
258    #[test]
259    fn test_count_than() {
260        let mut b = BinaryTrie::new();
261
262        for i in 0..1000 {
263            b.insert(i);
264            assert_eq!(b.count_less(i), i as u64);
265        }
266
267        assert_eq!(b.min().unwrap(), 0);
268        assert_eq!(b.max().unwrap(), 999);
269        for i in 0..1000 {
270            assert_eq!(b.count_more(i), 999 - i as u64);
271        }
272        assert_eq!(b.count_less(u32::MAX), 1000);
273        assert_eq!(b.count_more(u32::MIN), 999);
274    }
275
276    #[test]
277    fn library_checker() {
278        let mut b = BinaryTrie::new();
279        let query = [(0, 6), (0, 7), (2, 5), (1, 7), (1, 10), (2, 7)];
280        let mut ans = vec![];
281        query.iter().for_each(|&(p, x)| match p {
282            0 => {
283                b.insert(x);
284            }
285            1 => {
286                b.erase_all(x);
287            }
288            _ => ans.push(b.xor_min(x).unwrap_or_else(|| panic!("{}", x.to_string()))),
289        });
290
291        assert_eq!(vec![2, 1], ans);
292        assert_eq!(b.count(6), 1);
293        assert_eq!(b.count(7), 0);
294    }
295
296    #[test]
297    fn q() {
298        use crate::other::xorshift::XorShift;
299        let mut xs = XorShift::<u64>::new();
300        let mut b = BinaryTrie::new();
301        b.insert(0);
302        let mut ans = vec![];
303        for i in 0..200_000 {
304            match xs.next().unwrap() % 3 {
305                0 => {
306                    b.insert(xs.next().unwrap() as u32 % u32::MAX);
307                }
308                1 => {
309                    b.erase_all(xs.next().unwrap() as u32 % u32::MAX);
310                }
311                _ => ans.push(
312                    b.xor_min(xs.next().unwrap() as u32 % u32::MAX)
313                        .unwrap_or_else(|| panic!()),
314                ),
315            }
316
317            b.xor_min(i);
318        }
319    }
320    #[test]
321    fn lower_bound() {
322        let v = [
323            1, 1, 4, 7, 8, 9, 11, 64, 98, 641, 1_111, 1_111, 1_111, 6_000, 10_000, 123_456,
324            1_111_111, 9_999_999,
325        ];
326
327        let mut b = BinaryTrie::new();
328        v.iter().for_each(|x| {
329            b.insert(*x);
330        });
331        assert_eq!(b.lower_bound(0), None);
332        assert_eq!(b.lower_bound(1), Some(1));
333        assert_eq!(b.lower_bound(4), Some(4));
334        assert_eq!(b.lower_bound(7), Some(7));
335        assert_eq!(b.lower_bound(8), Some(8));
336        assert_eq!(b.lower_bound(9), Some(9));
337        assert_eq!(b.lower_bound(11), Some(11));
338        assert_eq!(b.lower_bound(64), Some(64));
339        assert_eq!(b.lower_bound(98), Some(98));
340        assert_eq!(b.lower_bound(641), Some(641));
341        assert_eq!(b.lower_bound(1_111), Some(1_111));
342        assert_eq!(b.lower_bound(6_000), Some(6_000));
343        assert_eq!(b.lower_bound(10_000), Some(10_000));
344        assert_eq!(b.lower_bound(123_456), Some(123_456));
345        assert_eq!(b.lower_bound(1_111_111), Some(1_111_111));
346        assert_eq!(b.lower_bound(9_999_999), Some(9_999_999));
347    }
348
349    #[test]
350    fn upper_bound() {
351        let v = [
352            1, 1, 4, 7, 8, 9, 11, 64, 98, 641, 1_111, 1_111, 1_111, 6_000, 10_000, 123_456,
353            1_111_111, 9_999_999,
354        ];
355
356        let mut b = BinaryTrie::new();
357        v.iter().for_each(|x| {
358            b.insert(*x);
359        });
360        assert_eq!(b.upper_bound(0), Some(1));
361        assert_eq!(b.upper_bound(1), Some(4));
362        assert_eq!(b.upper_bound(4), Some(7));
363        assert_eq!(b.upper_bound(7), Some(8));
364        assert_eq!(b.upper_bound(8), Some(9));
365        assert_eq!(b.upper_bound(9), Some(11));
366        assert_eq!(b.upper_bound(11), Some(64));
367        assert_eq!(b.upper_bound(64), Some(98));
368        assert_eq!(b.upper_bound(98), Some(641));
369        assert_eq!(b.upper_bound(641), Some(1_111));
370        assert_eq!(b.upper_bound(1_111), Some(6_000));
371        assert_eq!(b.upper_bound(6_000), Some(10_000));
372        assert_eq!(b.upper_bound(10_000), Some(123_456));
373        assert_eq!(b.upper_bound(123_456), Some(1_111_111));
374        assert_eq!(b.upper_bound(1_111_111), Some(9_999_999));
375        assert_eq!(b.upper_bound(9_999_999), None);
376    }
377
378    #[test]
379    fn c() {
380        let mut b = BinaryTrie::new();
381        b.insert(1);
382        assert_eq!(b.xth_element(0), None);
383        assert_eq!(b.xth_element(1), Some(1));
384    }
385}
competitive_library/structure/binary_trie.rs

competitive_library/structure/
binary_trie.rs
