[LeetCode] 208. Implement Trie (Prefix Tree)（实现字典树）

• Difficulty: Medium

• Related Topics: Design, Trie

• Link: https://leetcode.com/problems/implement-trie-prefix-tree/

Description

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

Example

Trie trie = new Trie();

trie.insert("apple");
trie.search("apple");   // returns true
trie.search("app");     // returns false
trie.startsWith("app"); // returns true
trie.insert("app");   
trie.search("app");     // returns true

Note

• You may assume that all inputs are consist of lowercase letters a-z.
• All inputs are guaranteed to be non-empty strings.

Solution

这题是设计题，设计一个字典树数据结构。相关的资料可以自行百度，这里给一个简单实现。代码如下（相关地方已做了注释）

class Trie() {

    /** Initialize your data structure here. */
    private val root = TrieNode()
    private class TrieNode {
        /** 经过该节点的单词个数（本题用不上，不过存储这个可以用来统计有多少个包含前缀的单词） */
        var wordSize: Int = 1
        /** 子节点。这里用哈希表实现，可以适应不同的字符 */
        val children = hashMapOf<Char, TrieNode>()
        /** 标记是否到达一个单词的结束 */
        var end = false
    }
    

    /** Inserts a word into the trie. */
    fun insert(word: String) {
        if (word.isEmpty()) {
            return
        }

        var node = root
        for (i in word.indices) {
            val c = word[i]
            // 如果这个字符子节点存在，增加其单词计数，否则就新建一个
            if (node.children.containsKey(c)) {
                node.children.getValue(c).wordSize++
            } else {
                node.children[c] = TrieNode()
            }

            // 已经遍历到末尾，给末尾节点增加结束标志
            if (i == word.lastIndex) {
                node.children[c]?.end = true
            }

            node = node.children.getValue(c)
        }
    }

    /** Returns if the word is in the trie. */
    fun search(word: String): Boolean {
        val node = getNode(word)
        return node?.end == true
    }

    /** Returns if there is any word in the trie that starts with the given prefix. */
    fun startsWith(prefix: String): Boolean {
    	return getNode(prefix) != null
    }

    /**
     * 获取到 `s` 末尾的节点，找不到则返回 `null`
     */
    private fun getNode(s: CharSequence): TrieNode? {
        if (s.isEmpty()) {
            return null
        }

        var node: TrieNode? = root
        for (c in s) {
            if (node?.children?.containsKey(c) != true) {
                return null
            }
            node = node.children[c]
        }
        return node
    }
}