14、Trie

内容来自刘宇波老师算法与数据结构体系课

1、Trie 的结构

2、实现 Trie

/**
 * 字典树: 非递归实现
 * 查询的时间复杂度为: O(w), w 为 word 的长度! 跟 Trie 中存储的 word 数量无关
 */
public class Trie {

    private class Node {
        public boolean isWord;
        public TreeMap<Character, Node> next;

        public Node(boolean isWord) {
            this.isWord = isWord;
            this.next = new TreeMap<>();
        }

        public Node() {
            this(false);
        }
    }

    private final Node root;
    private int size;

    public Trie() {
        root = new Node();
        size = 0;
    }

    public int getSize() {
        return size;
    }

    /**
     * 向 Trie 中添加一个新的单词 word
     */
    public void add(String word) {
        Node cur = root;

        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.next.containsKey(c)) cur.next.put(c, new Node());
            cur = cur.next.get(c);
        }

        if (!cur.isWord) {
            cur.isWord = true;
            size++;
        }
    }

    /**
     * 查询 Trie 中是否包含单词 word
     */
    public boolean contains(String word) {
        Node cur = root;

        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.next.containsKey(c)) return false;
            cur = cur.next.get(c);
        }

        return cur.isWord;
    }

    /**
     * 查询 Trie 中是否有以 prefix 为前缀的单词
     */
    public boolean isPrefix(String prefix) {
        Node cur = root;

        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (!cur.next.containsKey(c)) return false;
            cur = cur.next.get(c);
        }

        return true;
    }
}

3、添加与搜索单词 - 数据结构设计

211 - 添加与搜索单词 - 数据结构设计

/**
 * 模式匹配: 每个 '.' 都可以表示任何一个字母
 */
public class WordDictionary {

    private static class Node {
        public boolean isWord;
        public TreeMap<Character, Node> next;

        public Node(boolean isWord) {
            this.isWord = isWord;
            this.next = new TreeMap<>();
        }

        public Node() {
            this(false);
        }
    }

    private final Node root;

    public WordDictionary() {
        root = new Node();
    }

    public void addWord(String word) {
        Node cur = root;

        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.next.containsKey(c)) cur.next.put(c, new Node());
            cur = cur.next.get(c);
        }

        if (!cur.isWord) cur.isWord = true;
    }

    public boolean search(String word) {
        return match(root, word, 0);
    }

    /**
     * 在以 node 为根节点的 Trie 中查找 word[index ... n - 1]
     * node 中存放的是已经查询过的, 未查询的在 node 的下面
     */
    private boolean match(Node node, String word, int index) {
        if (index == word.length()) return node.isWord;
        
        char c = word.charAt(index);
        
        if (c != '.') {
            if (!node.next.containsKey(c)) return false;
            return match(node.next.get(c), word, index + 1);
        } else {
            for (Node nextNode : node.next.values()) {
                if (match(nextNode, word, index + 1)) return true;
            }
            return false;
        }
    }
}

4、键值映射

677 - 键值映射

public class MapSum {

    private static class Node {
        public int val;
        public TreeMap<Character, Node> next;

        public Node(int val) {
            this.val = val;
            next = new TreeMap<>();
        }

        public Node() {
            this(0);
        }
    }

    private final Node root;

    public MapSum() {
        root = new Node();
    }

    public void insert(String word, int val) {
        Node cur = root;

        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!cur.next.containsKey(c)) cur.next.put(c, new Node());
            cur = cur.next.get(c);
        }

        cur.val = val;
    }

    public int sum(String prefix) {
        Node cur = root;

        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (!cur.next.containsKey(c)) return 0;
            cur = cur.next.get(c);
        }

        return sum(cur);
    }

    private int sum(Node node) {
        int sum = node.val;
        if (node.next.isEmpty()) return sum;

        for (Node nextNode : node.next.values()) sum += sum(nextNode);
        return sum;
    }
}