use std::collections::BTreeSet;
pub(crate) const TRIE_ASCII_SIZE: usize = 256;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct TrieNode<const N: usize> {
end: bool, children: [Option<Box<TrieNode<N>>>; N],
}
impl<const N: usize> TrieNode<N> {
pub(crate) fn new() -> Self {
let children = std::array::from_fn(|_| None);
TrieNode {
end: false,
children,
}
}
pub(crate) fn create_end_node() -> Self {
let children = std::array::from_fn(|_| None);
TrieNode {
end: true,
children,
}
}
pub(crate) fn set_end_node(&mut self) {
self.end = true;
}
pub(crate) fn insert(&mut self, key: &str) {
let mut x = self;
let count = key.chars().count();
for (i, c) in key.chars().enumerate() {
let idx = (c as u8) as usize;
if idx > TRIE_ASCII_SIZE {
return;
}
if x.children[idx].is_none() {
if i == count - 1 {
x.children[idx] = Some(Box::new(TrieNode::create_end_node()));
} else {
x.children[idx] = Some(Box::new(TrieNode::new()));
}
} else if i == count - 1
&& let Some(child) = x.children[idx].as_deref_mut()
{
child.set_end_node();
}
#[allow(clippy::expect_used)]
let child = x.children[idx]
.as_deref_mut()
.expect("by prev if statement this value is always set");
x = child;
}
}
pub(crate) fn auto_complete(&self, prefix: &str) -> Option<Vec<String>> {
let mut x = self;
let mut res = vec![];
for c in prefix.chars() {
let i = (c as u8) as usize;
if i > TRIE_ASCII_SIZE {
return None;
}
x = x.children[i].as_deref()?;
}
let postfixs = x.dfs();
for postfix in postfixs {
let str = format!("{prefix}{postfix}");
res.push(str);
}
Some(res)
}
fn dfs(&self) -> BTreeSet<String> {
let mut q = vec![];
let mut res = BTreeSet::new();
let x = self;
let buf = String::new();
q.push((x, buf));
while let Some((node, curr_str)) = q.pop() {
for (i, child) in node.children.iter().enumerate() {
match child {
None => {}
Some(n) => {
if let Ok(val) = u8::try_from(i) {
let c = val as char;
let new_str = format!("{curr_str}{c}");
q.push((n.as_ref(), new_str));
}
}
}
}
if node.end {
res.insert(curr_str.clone());
}
}
res
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn auto_complete() {
let mut trie: TrieNode<TRIE_ASCII_SIZE> = TrieNode::new();
trie.insert("echo");
trie.insert("exit");
assert_eq!(
trie.auto_complete("e"),
Some(vec!["echo".into(), "exit".into()])
);
assert_eq!(trie.auto_complete("ech"), Some(vec!["echo".into()]));
}
#[test]
fn auto_complete_overlap() {
let mut trie: TrieNode<TRIE_ASCII_SIZE> = TrieNode::new();
trie.insert("xyz_dog");
trie.insert("xyz_dog_owl");
trie.insert("xyz_dog_owl_cow");
assert_eq!(
trie.auto_complete("xyz_"),
Some(vec![
"xyz_dog".into(),
"xyz_dog_owl".into(),
"xyz_dog_owl_cow".into()
])
);
}
}