qitoy_binary_trie/
lib.rs

1use std::num::NonZeroUsize;
2use std::ops::{BitXor, BitXorAssign, Shl, Shr};
3
4#[derive(Clone, Copy, Debug, Default)]
5struct TrieNode {
6    /// offset ^ digit == child index
7    offset: Option<NonZeroUsize>,
8    /// parrent index
9    parrent: usize,
10    /// subtree node count
11    cnt: usize,
12}
13
14#[derive(Debug, Default)]
15pub struct BinaryTrie<T> {
16    data: Vec<TrieNode>,
17    lazy: T,
18}
19
20impl<T: UInt> BinaryTrie<T> {
21    pub fn new() -> Self {
22        Self {
23            // [root, dummy]
24            data: vec![Default::default(); 2],
25            lazy: Default::default(),
26        }
27    }
28
29    pub fn len(&self) -> usize {
30        self.data[0].cnt
31    }
32
33    pub fn is_empty(&self) -> bool {
34        self.len() == 0
35    }
36
37    pub fn insert(&mut self, value: T) {
38        let value = value ^ self.lazy;
39        let mut idx = 0;
40        self.data[idx].cnt += 1;
41        for i in (0..T::BITS).rev() {
42            if self.data[idx].offset.is_none() {
43                let offset = self.push();
44                self.data[idx].offset = NonZeroUsize::new(offset);
45                self.data[offset].parrent = idx;
46                self.data[offset ^ 1].parrent = idx;
47            }
48            let offset = self.data[idx].offset.unwrap().get();
49            idx = offset | ((value >> i).as_usize() & 1);
50            self.data[idx].cnt += 1;
51        }
52    }
53
54    fn push(&mut self) -> usize {
55        let len = self.data.len();
56        self.data.resize(len + 2, Default::default());
57        len
58    }
59
60    pub fn remove(&mut self, value: T) {
61        if let Some(mut idx) = self.find(value) {
62            while idx != 0 {
63                self.data[idx].cnt -= 1;
64                idx = self.data[idx].parrent;
65            }
66            self.data[idx].cnt -= 1;
67        }
68    }
69
70    fn find(&self, value: T) -> Option<usize> {
71        let value = value ^ self.lazy;
72        let mut idx = 0;
73        for i in (0..T::BITS).rev() {
74            let offset = self.data[idx].offset?.get();
75            idx = offset | ((value >> i).as_usize() & 1);
76        }
77        (self.data[idx].cnt != 0).then_some(idx)
78    }
79
80    pub fn count(&self, value: T) -> usize {
81        self.find(value).map_or(0, |idx| self.data[idx].cnt)
82    }
83
84    pub fn xor_all(&mut self, value: T) {
85        self.lazy ^= value;
86    }
87
88    pub fn min(&self) -> Option<T> {
89        if self.is_empty() {
90            return None;
91        }
92        let mut ret = T::default();
93        let mut idx = 0;
94        for i in (0..T::BITS).rev() {
95            let offset = self.data[idx].offset.unwrap().get();
96            idx = offset ^ ((self.lazy >> i).as_usize() & 1);
97            if self.data[idx].cnt == 0 {
98                idx ^= 1;
99                ret ^= T::one() << i;
100            }
101        }
102        Some(ret)
103    }
104}
105
106pub trait UInt:
107    Default
108    + Copy
109    + BitXorAssign
110    + BitXor<Output = Self>
111    + Shl<u32, Output = Self>
112    + Shr<u32, Output = Self>
113{
114    const BITS: u32;
115    fn as_usize(self) -> usize;
116    fn one() -> Self;
117}
118
119impl UInt for u32 {
120    const BITS: u32 = u32::BITS;
121
122    fn as_usize(self) -> usize {
123        self as usize
124    }
125    fn one() -> Self {
126        1
127    }
128}
129
130impl UInt for u64 {
131    const BITS: u32 = u64::BITS;
132
133    fn as_usize(self) -> usize {
134        self as usize
135    }
136    fn one() -> Self {
137        1
138    }
139}
140
141impl UInt for usize {
142    const BITS: u32 = usize::BITS;
143
144    fn as_usize(self) -> usize {
145        self
146    }
147    fn one() -> Self {
148        1
149    }
150}
qitoy_binary_trie/lib.rs

qitoy_binary_trie/
lib.rs
