competitive_library/structure/
quaternary_trie.rs

1//! QuaternaryTrie
2
3#[derive(Debug, Default, Clone, PartialEq, Eq)]
4struct Node {
5    children: [Option<usize>; 4],
6    count: u64,
7}
8impl Node {
9    #[inline]
10    fn new() -> Self {
11        Self {
12            children: [None; 4],
13            count: 0,
14        }
15    }
16    #[inline]
17    fn get_child(&self, index: usize) -> &Option<usize> {
18        unsafe { self.children.get_unchecked(index) }
19    }
20    #[inline]
21    fn get_child_mut(&mut self, index: usize) -> &mut Option<usize> {
22        unsafe { self.children.get_unchecked_mut(index) }
23    }
24}
25
26#[derive(Debug, Default, Clone, PartialEq, Eq)]
27pub struct QuaternaryTrie {
28    nodes: Vec<Node>,
29    bit_length: u32,
30}
31impl QuaternaryTrie {
32    /// 構築
33    #[inline]
34    pub fn new() -> Self {
35        Self {
36            nodes: vec![Node::new()],
37            bit_length: 30,
38        }
39    }
40
41    #[inline]
42    fn get_node_mut(&mut self, index: usize) -> &mut Node {
43        unsafe { self.nodes.get_unchecked_mut(index) }
44    }
45    #[inline]
46    fn get_node(&self, index: usize) -> &Node {
47        unsafe { self.nodes.get_unchecked(index) }
48    }
49
50    /// 値の挿入
51    #[inline]
52    pub fn insert(&mut self, x: u32) -> u64 {
53        self.insert_n(x, 1)
54    }
55    #[inline]
56    pub fn insert_n(&mut self, x: u32, n: u64) -> u64 {
57        if n == 0 {
58            return 0;
59        }
60        let mut node_index = 0;
61        for i in (0..self.bit_length / 2).rev() {
62            self.get_node_mut(node_index).count += n;
63
64            node_index = match self
65                .get_node(node_index)
66                .get_child((x >> (i * 2) & 3) as usize)
67            {
68                Some(i) => *i,
69                None => {
70                    self.nodes.push(Node::new());
71                    *self
72                        .get_node_mut(node_index)
73                        .get_child_mut((x >> (i * 2) & 3) as usize) = Some(self.nodes.len() - 1);
74                    self.nodes.len() - 1
75                }
76            };
77        }
78        self.get_node_mut(node_index).count += n;
79        self.get_node(node_index).count
80    }
81
82    /// xのカウント
83    #[inline]
84    pub fn count(&self, x: u32) -> u64 {
85        let mut node_index = Some(0);
86
87        for i in (0..self.bit_length / 2).rev() {
88            if node_index.is_none() {
89                return 0;
90            }
91            node_index = *self
92                .get_node(node_index.unwrap())
93                .get_child((x >> (i * 2) & 3) as usize);
94        }
95        if node_index.is_none() {
96            return 0;
97        }
98        self.get_node(node_index.unwrap()).count
99    }
100
101    /// 値の削除
102    #[inline]
103    pub fn erase(&mut self, x: u32) -> Option<()> {
104        if 1 > self.count(x) {
105            return None;
106        }
107        self.inner_erase(x, 1)
108    }
109
110    /// 値をすべて削除
111    #[inline]
112    pub fn erase_all(&mut self, x: u32) -> Option<()> {
113        let erase_count = self.count(x);
114        if erase_count == 0 {
115            return None;
116        }
117        self.inner_erase(x, erase_count)
118    }
119
120    /// 値を削除
121    /// 内部関数
122    #[inline]
123    fn inner_erase(&mut self, x: u32, erase_count: u64) -> Option<()> {
124        let mut node_index = Some(0);
125        for i in (0..self.bit_length / 2).rev() {
126            self.get_node_mut(node_index?).count -= erase_count;
127            node_index = *self
128                .get_node(node_index?)
129                .get_child((x >> (i * 2) & 3) as usize);
130        }
131        self.get_node_mut(node_index?).count -= erase_count;
132
133        Some(())
134    }
135
136    /// xor 後の最小値を求める
137    #[inline]
138    pub fn xor_min(&self, x: u32) -> Option<u32> {
139        let mut ans = 0;
140
141        let mut node_index = Some(0);
142        for i in (0..self.bit_length / 2).rev() {
143            let bit = {
144                let mut buff = (x >> (i * 2) & 3) as usize;
145                let a = self.get_node(node_index.unwrap());
146
147                for j in 0..4 {
148                    if a.get_child(buff ^ j)
149                        .filter(|&index| self.get_node(index).count > 0)
150                        .is_some()
151                    {
152                        buff ^= j;
153                        break;
154                    }
155                }
156                buff
157            };
158            ans ^= (bit as u32) << (i * 2);
159            node_index = *self.get_node(node_index.unwrap()).get_child(bit);
160        }
161        Some(ans ^ x)
162    }
163
164    #[inline]
165    pub fn size(&self) -> u64 {
166        self.get_node(0).count
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn library_checker() {
176        let mut b = QuaternaryTrie::new();
177        let query = [(0, 6), (0, 7), (2, 5), (1, 7), (1, 10), (2, 7)];
178        let mut ans = vec![];
179        query.iter().for_each(|&(p, x)| match p {
180            0 => {
181                b.insert(x);
182            }
183            1 => {
184                b.erase_all(x);
185            }
186            _ => ans.push(b.xor_min(x).unwrap_or_else(|| panic!("{}", x.to_string()))),
187        });
188
189        assert_eq!(vec![2, 1], ans);
190        assert_eq!(b.count(6), 1);
191        assert_eq!(b.count(7), 0);
192    }
193}
competitive_library/structure/quaternary_trie.rs

competitive_library/structure/
quaternary_trie.rs
