add trie (#28)

* add trie

* fixes

* nit

* add type to doc

* trie default nits

Author: cameroncuster

Cargo.toml

@@ -30,6 +30,10 @@ path = "examples/data_structures/fenwick_yosupo.rs"
 name = "rmq"
 path = "examples/data_structures/rmq.rs"
 
+[[example]]
+name = "trie"
+path = "examples/data_structures/trie.rs"
+
 [[example]]
 name = "dijk_aizu"
 path = "examples/graphs/dijk_aizu.rs"

examples/data_structures/trie.rs

@@ -0,0 +1,32 @@
+// verification-helper: PROBLEM https://onlinejudge.u-aizu.ac.jp/courses/lesson/1/ALDS1/all/ALDS1_4_C
+
+use proconio::input;
+use programming_team_code_rust::data_structures::trie::Trie;
+
+fn main() {
+    input! {
+        n: usize,
+    }
+
+    let mut trie = Trie::default();
+    for _ in 0..n {
+        input! {
+            command: String,
+            key: String,
+        }
+
+        match command.as_str() {
+            "insert" => {
+                trie.insert(&key);
+            }
+            "find" => {
+                if trie.find(&key) > 0 {
+                    println!("yes");
+                } else {
+                    println!("no");
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
+}

src/data_structures/mod.rs

@@ -2,3 +2,4 @@
 pub mod dsu;
 pub mod fenwick;
 pub mod rmq;
+pub mod trie;

src/data_structures/trie.rs

@@ -0,0 +1,79 @@
+//! # Trie
+
+const ALPHABET_SIZE: usize = 26;
+const FIRST_CHAR: char = 'A';
+
+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;
+///
+/// let mut trie = Trie::default();
+/// trie.insert("HELLO");
+/// trie.insert("HELLO");
+/// trie.insert("WORLD");
+/// assert_eq!(trie.find("HELLO"), 2);
+/// assert_eq!(trie.find("WORLD"), 1);
+/// assert_eq!(trie.find("WORLDS"), 0);
+/// ```
+pub struct Trie {
+    t: Vec<Node>,
+}
+
+impl Default for Trie {
+    fn default() -> Self {
+        Trie {
+            t: vec![Node::default()],
+        }
+    }
+}
+
+impl Trie {
+    /// Insert a string into the Trie
+    ///
+    /// # Complexity
+    /// - Time: O(|s|)
+    /// - Space: O(|s|)
+    pub fn insert(&mut self, s: &str) {
+        let mut v = 0;
+        for ch in s.chars() {
+            let idx = (ch as u8 - FIRST_CHAR as u8) as usize;
+            if self.t[v].next[idx].is_none() {
+                self.t[v].next[idx] = Some(self.t.len());
+                self.t.push(Node::default());
+            }
+            v = self.t[v].next[idx].unwrap();
+        }
+        self.t[v].cnt_words += 1;
+    }
+
+    /// Find the number of times a string appears in the Trie
+    ///
+    /// # Complexity
+    /// - Time: O(|s|)
+    /// - Space: O(1)
+    pub fn find(&self, s: &str) -> usize {
+        let mut v = 0;
+        for ch in s.chars() {
+            let idx = (ch as u8 - FIRST_CHAR as u8) as usize;
+            if self.t[v].next[idx].is_none() {
+                return 0;
+            }
+            v = self.t[v].next[idx].unwrap();
+        }
+        self.t[v].cnt_words
+    }
+}