trie树

  1. use std::collections::HashMap;
  2. use std::hash::Hash;
  4. #[derive(Debug, Default)]
  5. struct Node<Key: Default, Type: Default> {
  6.     children: HashMap<Key, Node<Key, Type>>,
  7.     value: Option<Type>,
  8. }
  10. #[derive(Debug, Default)]
  11. pub struct Trie<Key, Type>
  12. where
  13.     Key: Default + Eq + Hash,
  14.     Type: Default,
  15. {
  16.     root: Node<Key, Type>,
  17. }
  19. impl<Key, Type> Trie<Key, Type>
  20. where
  21.     Key: Default + Eq + Hash,
  22.     Type: Default,
  23. {
  24.     pub fn new() -> Self {
  25.         Self {
  26.             root: Node::default(),
  27.         }
  28.     }
  30.     pub fn insert(&mut self, key: impl IntoIterator<Item = Key>, value: Type)
  31.     where
  32.         Key: Eq + Hash,
  33.     {
  34.         let mut node = &mut self.root;
  35.         for c in key.into_iter() {
  36.             node = node.children.entry(c).or_insert_with(Node::default);
  37.         }
  38.         node.value = Some(value);
  39.     }
  41.     pub fn get(&self, key: impl IntoIterator<Item = Key>) -> Option<&Type>
  42.     where
  43.         Key: Eq + Hash,
  44.     {
  45.         let mut node = &self.root;
  46.         for c in key.into_iter() {
  47.             if node.children.contains_key(&c) {
  48.                 node = node.children.get(&c).unwrap()
  49.             } else {
  50.                 return None;
  51.             }
  52.         }
  53.         node.value.as_ref()
  54.     }
  55. }
  57. #[cfg(test)]
  58. mod tests {
  60.     use super::*;
  62.     #[test]
  63.     fn test_insertion() {
  64.         let mut trie = Trie::new();
  65.         assert_eq!(trie.get("".chars()), None);
  67.         trie.insert("foo".chars(), 1);
  68.         trie.insert("foobar".chars(), 2);
  70.         let mut trie = Trie::new();
  71.         assert_eq!(trie.get(vec![1, 2, 3]), None);
  73.         trie.insert(vec![1, 2, 3], 1);
  74.         trie.insert(vec![3, 4, 5], 2);
  75.     }
  77.     #[test]
  78.     fn test_get() {
  79.         let mut trie = Trie::new();
  80.         trie.insert("foo".chars(), 1);
  81.         trie.insert("foobar".chars(), 2);
  82.         trie.insert("bar".chars(), 3);
  83.         trie.insert("baz".chars(), 4);
  85.         assert_eq!(trie.get("foo".chars()), Some(&1));
  86.         assert_eq!(trie.get("food".chars()), None);
  88.         let mut trie = Trie::new();
  89.         trie.insert(vec![1, 2, 3, 4], 1);
  90.         trie.insert(vec![42], 2);
  91.         trie.insert(vec![42, 6, 1000], 3);
  92.         trie.insert(vec![1, 2, 4, 16, 32], 4);
  94.         assert_eq!(trie.get(vec![42, 6, 1000]), Some(&3));
  95.         assert_eq!(trie.get(vec![43, 44, 45]), None);
  96.     }
  97. }