[Leetcode] Implement Trie 实现前缀树

Implement Trie

Implement a trie with insert, search, and startsWith methods.

Note: You may assume that all inputs are consist of lowercase letters a-z.

哈希表法

复杂度

时间 插入和查询都是O(K) K是词的长度 空间 O(NK) N是字典里词的个数

思路

前缀树的具体讲解请戳这篇博客。这里我们实现树节点时使用了哈希表来映射字母和子节点的关系。
insert()：对于插入操作，我们遍历字符串同时，根据上一个节点的哈希表来找到下一个节点，直到哈希表中没有相应的字母，我们就新建一个节点。然后从这个新建节点开始，用同样的方法把剩余的字母添加完。记住最后一个字母要添加叶子节点的标记，表明这个词到此已经完整了。
search()：对于搜索操作，我们也是遍历字符串，然后根据每个节点的哈希表找到路径，最后返回该字符串最后一个字母所在节点。如果中途有找不到路径的情况就直接返回null，如果找到了最后的节点，如果它也是叶子结点的话，就说明找到了。
startWith()：使用和search()，一样的方法，只是我们返回的节点不用判断是否是叶子节点。只要找到就行。

代码

class TrieNode {
    // Initialize your data structure here.
    HashMap<Character, TrieNode> children = new HashMap<Character, TrieNode>();
    boolean isLeaf = false;
    char c;
    public TrieNode(){}
    public TrieNode(char c) {
        this.c = c;
    }
}

public class Trie {
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    // Inserts a word into the trie.
    public void insert(String word) {
        HashMap<Character, TrieNode> children = root.children;
        for(int i = 0; i < word.length(); i++){
            TrieNode next;
            // 如果已有该字母的节点，则转向该节点
            if(children.containsKey(word.charAt(i))){
                next = children.get(word.charAt(i));
            } else {
            // 如果没有该字母的节点，就新建一个节点
                next = new TrieNode(word.charAt(i));
                children.put(word.charAt(i), next);
            }
            children = next.children;
            if(i == word.length() - 1){
                next.isLeaf = true;
            }
        }
    }

    // Returns if the word is in the trie.
    public boolean search(String word) {
        TrieNode res = searchNode(word);
        if(res != null && res.isLeaf){
            return true;
        } else {
            return false;
        }
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    public boolean startsWith(String prefix) {
        return searchNode(prefix) != null;
    }
    
    private TrieNode searchNode(String word){
        HashMap<Character, TrieNode> children = root.children;
        TrieNode next = null;
        for(int i = 0; i < word.length(); i++){
            if(children.containsKey(word.charAt(i))){
                next = children.get(word.charAt(i));
                children = next.children;
            } else {
                return null;
            }
        }
        return next;
    }
}

// Your Trie object will be instantiated and called as such:
// Trie trie = new Trie();
// trie.insert("somestring");
// trie.search("key");

后续 Follow Up

Q：给定一个标准前缀树，请写一段程序将其压缩。
A：压缩前缀树其实就是将所有只有一个子节点的节点合并成一个，以减少没有意义的类似链表式的链接。

首先我们先将TrieNode稍微改一下。让它能存字符串而不只是字母。

class TrieNode {
    // Initialize your data structure here.
    HashMap<Character, TrieNode> children = new HashMap<Character, TrieNode>();
    boolean isLeaf = false;
    String str;
    public TrieNode(){}
    public TrieNode(char c) {
        this.str = String.valueOf(c);
    }
}

然后我们开始遍历这个前缀树。

public void compressTrie(Trie t){
    compress(t.getRoot());
}

private void compress(TrieNode n){
    if(n == null) return;
    if(n.children.size()==1){
        TrieNode next =    n.children.get(n.children.keySet().iterator().next());
        n.str = n.str + next.str;
        n.children = next.children;
        compress(next);
    } else {
        for(String key: n.children.keySet()){
            compress(n.children.get(key));
        }
    }
}