数据结构与算法(Java) 42：前缀树

package algorithm.section7;

public class TrieTree {
    public static class TrieNode{
        public int path;  // 存放所有前缀经过该节点的次数
        public int end;  //存放以该节点为最后一个值的前缀数
        public TrieNode[] next;  // 以26个字母为例：记录哪个分支存放了节点

        public TrieNode(){
            path = 0;
            end = 0;
            next = new TrieNode[26];
        }
    }

    public static class Trie{
        private TrieNode head;
        public Trie(){
            head = new TrieNode();
        }

        // 插入前缀：
        public void insert(String str){
            if (str.isEmpty()) return;

            char[] c = str.toCharArray();
            TrieNode flag = head;
            for (int i = 0; i < c.length; i++) {
                if (flag.next[c[i] - 'a'] == null)
                    flag.next[c[i] - 'a'] = new TrieNode();
                flag.path++;
                if (i == c.length - 1) flag.end++;
                flag = flag.next[c[i] - 'a'];
            }
        }

        // 删除前缀
        public void delete(String str){
            if (str.isEmpty()) return;

            char[] c = str.toCharArray();
            TrieNode flag = head;
            for (int i = 0; i < c.length; i++) {
                if (i == c.length - 1) flag.end--;
                if (flag.next[c[i] - 'a'].path > 1) {
                    flag.next[c[i] - 'a'].path--;
                    flag = flag.next[c[i] - 'a'];
                } else {
                    flag.next[c[i] - 'a'].path--;
                    flag.next[c[i] - 'a'] = null;
                    return;
                }
            }
        }

        // 查找某前缀插入过多少次
        public int search(String str){
            if (str.isEmpty()) return 0;

            char[] c = str.toCharArray();
            TrieNode flag = head;
            for (int i = 0; i < c.length; i++) {
                if (i == c.length - 1) return flag.end;
                if (flag.next[c[i] - 'a'] != null)
                    flag = flag.next[c[i] - 'a'];
                else return 0;
            }
            return 0;
        }

        // 查找有多少前缀以给定字符串为前缀
        public int prefixNumber(String str){
            if (str.isEmpty()) return 0;

            char[] c = str.toCharArray();
            TrieNode flag = head;

            for (int i = 0; i < c.length; i++){
                if (flag.next[c[i] - 'a'] != null)
                    flag = flag.next[c[i] - 'a'];
                else return 0;
            }
            return flag.path;
        }
    }

    public static void main(String[] args){
        Trie trie = new Trie();
        System.out.println(trie.search("zuo"));  // 0
        trie.insert("zuo");
        System.out.println(trie.search("zuo"));  // 1
        trie.delete("zuo");
        System.out.println(trie.search("zuo"));  // 0
        trie.insert("zuo");
        trie.insert("zuo");
        trie.delete("zuo");
        System.out.println(trie.search("zuo"));  // 1
        trie.delete("zuo");
        System.out.println(trie.search("zuo"));  // 0
        trie.insert("zuoa");
        trie.insert("zuoac");
        trie.insert("zuoab");
        trie.insert("zuoad");
        trie.delete("zuoa");
        System.out.println(trie.search("zuoa"));  // 0
        System.out.println(trie.prefixNumber("zuo"));  // 3
    }
}