Binary trie (#36)

* add binary trie

* small fix

* udpate binary trie template

* binary trie without template :)

* clippy

Author: cameroncuster

Cargo.toml

@@ -10,6 +10,10 @@ edition = "2024"
 [dependencies]
 proconio = "0.4.5"
 
+[[example]]
+name = "binary_trie"
+path = "examples/data_structures/binary_trie.rs"
+
 [[example]]
 name = "dsu_aizu"
 path = "examples/data_structures/dsu_aizu.rs"

examples/data_structures/binary_trie.rs

@@ -0,0 +1,35 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/set_xor_min
+
+use proconio::input;
+use programming_team_code_rust::data_structures::binary_trie::BinaryTrie;
+
+fn main() {
+    input! {
+        q: usize,
+    }
+
+    let mut trie = BinaryTrie::default();
+    for _ in 0..q {
+        input! {
+            t: u8,
+            x: usize,
+        }
+
+        match t {
+            0 => {
+                if trie.count(x) == 0 {
+                    trie.update(x, 1);
+                }
+            }
+            1 => {
+                if trie.count(x) == 1 {
+                    trie.update(x, -1);
+                }
+            }
+            2 => {
+                println!("{}", trie.min_xor(x));
+            }
+            _ => unreachable!(),
+        }
+    }
+}

src/data_structures/binary_trie.rs

@@ -0,0 +1,104 @@
+//! # Binary Trie which can be used as a multiset of integers
+
+#[derive(Default)]
+struct Node {
+    next: [Option<usize>; 2],
+    sub_sz: isize,
+}
+
+/// # Example
+/// ```
+/// use programming_team_code_rust::data_structures::binary_trie::BinaryTrie;
+///
+/// let mut trie = BinaryTrie::default();
+/// trie.update(1, 1);
+/// trie.update(2, 1);
+/// trie.update(3, 1);
+/// trie.update(2, -1);
+/// assert_eq!(trie.count(2), 0);
+/// assert_eq!(trie.count(3), 1);
+/// assert_eq!(trie.count(4), 0);
+/// assert_eq!(trie.min_xor(0), 1);
+/// assert_eq!(trie.min_xor(1), 0);
+/// assert_eq!(trie.min_xor(2), 1);
+/// assert_eq!(trie.min_xor(3), 0);
+/// assert_eq!(trie.min_xor(4), 5);
+/// ```
+pub struct BinaryTrie {
+    t: Vec<Node>,
+}
+
+impl Default for BinaryTrie {
+    fn default() -> Self {
+        BinaryTrie {
+            t: vec![Node::default()],
+        }
+    }
+}
+
+impl BinaryTrie {
+    /// Change the number of occurrences of `num` by `delta`
+    ///
+    /// # Complexity
+    /// - Time: O(log(max_num))
+    /// - Space: O(log(max_num))
+    pub fn update(&mut self, num: usize, delta: isize) {
+        let mut v = 0;
+        for i in (0..usize::BITS).rev() {
+            let bit = (num >> i) & 1;
+            if self.t[v].next[bit].is_none() {
+                self.t[v].next[bit] = Some(self.t.len());
+                self.t.push(Node::default());
+            }
+            self.t[v].sub_sz += delta;
+            v = self.t[v].next[bit].unwrap();
+        }
+        self.t[v].sub_sz += delta;
+    }
+
+    /// Count the number of occurrences of `num`
+    ///
+    /// # Complexity
+    /// - Time: O(log(max_num))
+    /// - Space: O(1)
+    pub fn count(&self, num: usize) -> isize {
+        let mut v = 0;
+        for i in (0..usize::BITS).rev() {
+            let bit = (num >> i) & 1;
+            if self.t[v].next[bit].is_none() {
+                return 0;
+            }
+            v = self.t[v].next[bit].unwrap();
+        }
+        self.t[v].sub_sz
+    }
+
+    /// Find the minimum xor of `num` and any number in the trie
+    ///
+    /// # Panics
+    /// ```panic
+    /// use programming_team_code_rust::data_structures::binary_trie::BinaryTrie;
+    ///
+    /// let trie = BinaryTrie::default();
+    /// trie.min_xor(0);
+    /// ```
+    ///
+    /// # Complexity
+    /// - Time: O(log(max_num))
+    /// - Space: O(1)
+    pub fn min_xor(&self, num: usize) -> usize {
+        assert!(self.t.len() > 1);
+        let mut v = 0;
+        let mut ans = 0;
+        for i in (0..usize::BITS).rev() {
+            let bit = (num >> i) & 1;
+            if self.t[v].next[bit].is_some() && self.t[self.t[v].next[bit].unwrap()].sub_sz > 0 {
+                v = self.t[v].next[bit].unwrap();
+            } else {
+                ans |= 1 << i;
+                v = self.t[v].next[bit ^ 1].unwrap();
+            }
+        }
+        ans
+    }
+}

src/data_structures/mod.rs

@@ -1,4 +1,5 @@
 //! # Data Structures
+pub mod binary_trie;
 pub mod dsu;
 pub mod fenwick;
 pub mod rmq;

src/data_structures/trie.rs

@@ -3,20 +3,12 @@
 const ALPHABET_SIZE: usize = 26;
 const FIRST_CHAR: char = 'A';
 
+#[derive(Default)]
 struct Node {
     next: [Option<usize>; ALPHABET_SIZE],
     cnt_words: usize,
 }
 
-impl Default for Node {
-    fn default() -> Self {
-        Node {
-            next: [None; ALPHABET_SIZE],
-            cnt_words: 0,
-        }
-    }
-}
-
 /// # Example
 /// ```
 /// use programming_team_code_rust::data_structures::trie::Trie;