Rust_leetcode 208. 实现 Trie (前缀树)

字典树用于字符串处理

代码实现

#![allow(unused)]

use std::collections::HashMap;
use std::cell::RefCell;
use std::rc::Rc;


// 节点定义
struct Node {
    is_word: bool,
    next: HashMap<char, Rc<RefCell<Node>>>,
}

impl Node {
    pub fn new_with_is_word(is_word: bool) -> Self {
        return Self {
            is_word,
            next: HashMap::new(),
        };
    }
    pub fn new() -> Self {
        return Self {
            is_word: false,
            next: HashMap::new(),
        };
    }
}

// 字典树
struct Trie {
    root: Rc<RefCell<Node>>,
    size: usize,
}

impl Trie {
    pub fn new() -> Self {
        return Trie {
            root: Rc::new(RefCell::new(Node::new())),
            size: 0,
        };
    }

    pub fn get_size(&self) -> usize {
        return self.size;
    }

    // 添加单词
    pub fn insert(&mut self, word: String) {
        let mut tmp = self.root.clone();
        // 遍历单词
        for c in word.chars() {
            let n = tmp.borrow().next.get(&c).cloned();
            if let Some(no) = n {
                tmp = no.clone();
            } else {
                let t = Rc::new(RefCell::new(Node::new()));
                tmp.borrow_mut().next.insert(c, t.clone());
                tmp = t.clone()
            };
        }
        if tmp.borrow().is_word == false {
            tmp.borrow_mut().is_word = true;
            self.size += 1;
        }
    }

    // 包含
    pub fn search(&self, word: String) -> bool {
        let mut tmp = self.root.clone();

        for c in word.chars() {
            let n = tmp.borrow().next.get(&c).cloned();
            if let Some(no) = n {
                tmp = no.clone();
            } else {
                return false;
            };
        }
        if tmp.borrow().is_word == true {
            return true;
        }
        false
    }

    pub fn starts_with(&self, word: String) -> bool {
        let mut tmp = self.root.clone();

        for c in word.chars() {
            let n = tmp.borrow().next.get(&c).cloned();
            if let Some(no) = n {
                tmp = no.clone();
            } else {
                return false;
            };
        }
        return true;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_works() {
        let mut trie = Trie::new();
        trie.insert("pan".to_string());
        trie.insert("panda".to_string());
        assert_eq!(trie.search("pan".to_string()), true);
        assert_eq!(trie.search("pa".to_string()), false);
        assert_eq!(trie.starts_with("pand".to_string()), true);
    }
}

结果