• 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 lettersa-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
    }
}