const R: usize = 256;
const DOT: usize = '.' as usize;
type Link<T> = Option<Box<Node<T>>>;
pub struct Node<T> {
pub val: Option<T>,
next: Vec<Link<T>>,
}
impl<T> Node<T> {
pub fn new() -> Link<T> {
let mut next = Vec::with_capacity(R);
for _ in 0..R {
next.push(None);
}
Some(Box::new(Node { val: None, next }))
}
}
#[derive(Default)]
pub struct TrieTS<T> {
root: Link<T>,
}
impl<T> TrieTS<T> {
pub fn get<'a>(&'a self, key: &str) -> Option<&'a Node<T>> {
Self::get_in(&self.root, key, 0)
}
fn get_in<'a>(root: &'a Link<T>, key: &str, d: usize) -> Option<&'a Node<T>> {
root.as_deref().and_then(|root| {
if d == key.len() {
Some(root)
} else {
Self::get_in(&root.next[Self::char_at(key, d)], key, d + 1)
}
})
}
fn char_at(s: &str, d: usize) -> usize {
s.as_bytes()[d] as usize
}
pub fn put(&mut self, key: &str, val: T) {
self.root = Self::put_in(self.root.take(), key, val, 0)
}
pub fn put_in(root: Link<T>, key: &str, val: T, d: usize) -> Link<T> {
let mut ans = if root.is_none() { Node::new() } else { root };
if d == key.len() {
ans.as_deref_mut().unwrap().val = Some(val);
} else {
let p = Self::char_at(key, d);
let node = &mut ans.as_deref_mut().unwrap().next[p];
let next = node.take();
*node = Self::put_in(next, key, val, d + 1)
};
ans
}
pub fn keys(&self) -> Vec<String> {
self.keys_with_prefix("")
}
pub fn keys_with_prefix(&self, prefix: &str) -> Vec<String> {
let mut q = vec![];
Self::collect(
Self::get_in(&self.root, prefix, 0),
prefix.to_string(),
&mut q,
);
q
}
pub fn keys_that_match(&self, pat: &str) -> Vec<String> {
let mut q = vec![];
Self::collect_by_match(self.root.as_deref(), "".to_string(), pat, &mut q);
q
}
fn collect_by_match(root: Option<&Node<T>>, prefix: String, pat: &str, q: &mut Vec<String>) {
if let Some(root) = root {
let d = prefix.len();
if d == pat.len() {
if root.val.is_some() {
q.push(prefix);
}
return;
}
let next_char = Self::char_at(pat, d);
for (i, next) in root
.next
.iter()
.enumerate()
.filter(|(i, _next)| DOT == next_char || *i == next_char)
{
Self::collect_by_match(
next.as_deref(),
format!("{}{}", prefix, char::from(i as u8)),
pat,
q,
)
}
}
}
fn collect(root: Option<&Node<T>>, prefix: String, q: &mut Vec<String>) {
if let Some(root) = root {
if root.val.is_some() {
q.push(prefix.to_string());
}
for (i, next) in root.next.iter().enumerate() {
Self::collect(
next.as_deref(),
format!("{}{}", prefix, char::from(i as u8)),
q,
)
}
}
}
pub fn longest_prefix_of(&self, s: &str) -> String {
let len = Self::search(self.root.as_deref(), s, 0, 0);
s[0..len].to_string()
}
fn search(root: Option<&Node<T>>, s: &str, d: usize, len: usize) -> usize {
root.map(|node| {
let mut len = len;
if node.val.is_some() {
len = d;
}
if d == s.len() {
len
} else {
let c = Self::char_at(s, d);
Self::search(node.next[c].as_deref(), s, d + 1, len)
}
})
.unwrap_or(len)
}
pub fn delete(&mut self, key: &str) {
self.root = Self::delete_in(self.root.take(), key, 0);
}
fn delete_in(root: Link<T>, key: &str, d: usize) -> Link<T> {
root.and_then(|mut root| {
if d == key.len() {
root.val.take();
} else {
let c = Self::char_at(key, d);
let next = root.next[c].take();
root.next[c] = Self::delete_in(next, key, d + 1);
}
if root.val.is_some() {
Some(root)
} else {
let has_child = root.next.iter().any(|f| f.is_some());
if has_child {
Some(root)
} else {
None
}
}
})
}
}
#[cfg(test)]
mod test {
use super::TrieTS;
#[test]
fn test() {
let mut st: TrieTS<Option<Option<usize>>> = TrieTS::default();
st.put("by", None);
st.put("sea", None);
st.put("sells", None);
st.put("she", None);
st.put("shells", None);
st.put("shore", None);
st.put("the", None);
println!("{:?}", st.keys());
assert!(st.keys_with_prefix("b").contains(&"by".to_string()));
let keys = st.keys_with_prefix("s");
assert!(keys.contains(&"sea".to_string()));
assert!(keys.contains(&"she".to_string()));
assert!(keys.len() == 5);
let keys = st.keys_with_prefix("shells");
assert!(keys.contains(&"shells".to_string()));
assert!(st.get("the").unwrap().val.is_some());
assert!(st.get("b").unwrap().val.is_none());
assert_eq!(st.keys_that_match("by").len(), 1);
assert_eq!(st.keys_that_match("b.").len(), 1);
assert_eq!(st.keys_that_match("s..").len(), 2);
assert_eq!(st.keys_that_match("s....").len(), 2);
assert_eq!(st.longest_prefix_of("shellsaaa"), "shells".to_string());
assert_eq!(st.longest_prefix_of("by"), "by".to_string());
assert_eq!(st.longest_prefix_of("b"), "".to_string());
st.delete("shells");
assert!(st.keys_that_match("shells").is_empty());
assert!(!st.keys().is_empty());
let mut len = st.keys().len();
for key in st.keys() {
st.delete(&key);
len -= 1;
assert_eq!(st.keys().len(), len);
}
assert!(st.keys().is_empty());
}
}