use core::borrow::Borrow;
use core::hash::Hash;
use std::collections::HashMap;
/// Trie (Prefix tree) data structure
#[derive(Debug)]
pub struct Trie<K: Hash + Eq, V = ()> {
root: TrieNode<K, V>,
}
impl<K: Hash + Eq, V> Trie<K, V> {
/// Creates a new Trie
pub fn new() -> Self {
Self { root: TrieNode::default() }
}
/// Inserts a key-value pair on the trie
///
/// # Example
/// ```
/// use prfx::Trie;
///
/// let mut trie = Trie::<char, String>::new();
/// trie.insert("hello".chars(), String::from("world!"));
/// trie.insert("hell".chars(), String::from("is hot!"));
/// ```
pub fn insert<I>(&mut self, key: I, val: V)
where
I: IntoIterator<Item = K>
{
self.root.insert(key, val);
}
/// Searches an entry inside the trie
///
/// # Example
/// ```
/// use prfx::Trie;
///
/// let mut trie = Trie::<char, String>::new();
/// trie.insert("hello".chars(), String::from("world!"));
/// trie.insert("hell".chars(), String::from("is hot!"));
///
/// assert_eq!(
/// trie.search("hello".chars()).map(String::as_str),
/// Some("world!")
/// );
/// assert_eq!(
/// trie.search("hell".chars()).map(String::as_str),
/// Some("is hot!")
/// );
/// assert_eq!(trie.search("he".chars()), None);
/// ```
pub fn search<I, Q>(&self, elems: I) -> Option<&V>
where
I: IntoIterator<Item = Q>,
K: Borrow<Q>,
Q: Eq + Hash,
{
self.root.search(elems)
}
/// Removes an entry from the trie
///
/// # Example
/// ```
/// use prfx::Trie;
///
/// let mut trie = Trie::<char, String>::new();
/// trie.insert("hello".chars(), String::from("world!"));
///
/// assert!(trie.search("hello".chars()).is_some());
/// trie.remove("hello".chars());
/// assert!(trie.search("hello".chars()).is_none());
/// ```
pub fn remove<I, Q>(&mut self, elems: I)
where
I: IntoIterator<Item = Q>,
K: Borrow<Q>,
Q: Eq + Hash,
{
let mut it = elems.into_iter();
let Some(first) = it.next() else { return };
let Some(node) = self.root.childs.get_mut(&first) else { return };
if node.remove_rec(&mut it) {
self.root.childs.remove(&first);
}
}
}
impl<K: Hash + Eq, V> Default for Trie<K, V> {
fn default() -> Self {
Self { root: Default::default() }
}
}
impl<K: Hash + Eq + Clone, V: Clone> Clone for Trie<K, V> {
fn clone(&self) -> Self {
Self { root: self.root.clone() }
}
}
#[derive(Debug)]
struct TrieNode<K: Hash + Eq, V> {
childs: HashMap<K, TrieNode<K, V>>,
value: Option<V>,
}
impl<K: Hash + Eq, V> TrieNode<K, V> {
fn insert<I>(&mut self, key: I, val: V)
where
I: IntoIterator<Item = K>
{
let mut current = self;
for elem in key {
current = current.childs.entry(elem).or_default();
}
current.value = Some(val);
}
fn search<I, Q>(&self, elems: I) -> Option<&V>
where
I: IntoIterator<Item = Q>,
K: Borrow<Q>,
Q: Eq + Hash,
{
let mut current = self;
for elem in elems {
current = current.childs.get(&elem)?;
}
current.value.as_ref()
}
/* Returns true if the caller should remove the `self` node */
fn remove_rec<I, Q>(&mut self, elems: &mut I) -> bool
where
I: Iterator<Item = Q>,
K: Borrow<Q>,
Q: Eq + Hash,
{
if let Some(next) = elems.next() {
let Some(child) = self.childs.get_mut(&next) else {
return false
};
if child.remove_rec(elems) {
self.childs.remove(&next);
}
} else {
self.value = None;
}
self.childs.is_empty() && self.value.is_none()
}
}
impl<K: Hash + Eq + Clone, V: Clone> Clone for TrieNode<K, V> {
fn clone(&self) -> Self {
Self { childs: self.childs.clone(), value: self.value.clone() }
}
}
impl<K: Hash + Eq, V> Default for TrieNode<K, V> {
fn default() -> Self {
Self { childs: Default::default(), value: None }
}
}
/// Useful functions for tries of `char`
impl Trie<char> {
#[inline]
pub fn insert_str(&mut self, s: &str) {
self.insert(s.chars(), ());
}
#[inline]
pub fn search_str(&self, elems: &str) -> bool {
self.search(elems.chars()).is_some()
}
#[inline]
pub fn remove_str(&mut self, elems: &str) {
self.remove(elems.chars());
}
}
/// Useful functions for tries without values
impl<K: Hash + Eq> Trie<K, ()> {
#[inline]
pub fn insert_key<I>(&mut self, elems: I)
where
I: IntoIterator<Item = K>
{
self.insert(elems, ());
}
#[inline]
pub fn search_key<I, Q>(&self, elems: I) -> bool
where
I: Iterator<Item = Q>,
K: Borrow<Q>,
Q: Eq + Hash,
{
self.search(elems).is_some()
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test() {
let mut tn: Trie<char> = Trie::default();
tn.insert_str("Hola");
tn.insert_str("Holgado");
assert!(tn.search_str("Hola"));
assert!(tn.search_str("Holgado"));
assert!(!tn.search_str("Hol"));
assert!(!tn.search_str("Holistico"));
tn.remove_str("Hola");
assert!(!tn.search_str("Hola"));
}
#[test]
fn spanish() {
const WORDS: &str = "Árbol,Casa,Gato,Perro,Cielo,Mar,Montaña,Libro,Sol,Luna,Río,Flor,Mesa,Silla,Ventana,Puerta,Niña,Carro,Camino,Bosque,Nube,Tiempo,Reloj,Playa,Lago,Pájaro,Estrella,Nieve,Viento,Fuego,Coche,Tren,Avión,Isla,Campo,Niño,Hombre,Mujer,Mano,Pie,Cabeza,Boca,Ojo,Oreja,Corazón,Amigo,Trabajo,Escuela,Maestro,Alumno,Música,Canción,Baile,Juego,Pelota,Fruta,Verdura,Pan,Agua,Vaso,Plato,Cuchara,Taza,Ropa,Zapato,Sombrero,Bolsa,Llave,Teléfono,Computadora,Foto,Carta,Papel,Revista,Periódico,Ciudad,Pueblo,Calle,Plaza,Parque,Jardín,Animal,León,Tigre,Elefante,Mono,Pescado,Mariposa,Serpiente,Hormiga,Araña,Ratón,Guitarra,Piano,Trompeta,Violín";
let mut dict = Trie::<char>::new();
let mut words: Vec<&str> = WORDS.split(',').collect();
for word in &words {
dict.insert_str(word);
}
words.sort();
for word in &words {
dbg!(word);
assert!(dict.search_str(word));
dict.remove_str(word);
assert!(!dict.search_str(word));
}
assert!(!dict.search_str("NOT_FOUND"));
}
}