HàPhan 河

Implement Trie (Prefix Tree)

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

Each node in Trie keeps 3 information: value, list of its child-nodes, and is it a leaf node?
Should use hashtable to keep children.

``````
class TrieNode<T: Hashable> {
var value: T?

weak var parent: TrieNode?
var children = [T: TrieNode]()

var isTerminate: Bool = false

init(_ value: T? = nil, _ parent: TrieNode? = nil) {
self.value = value
self.parent = parent
}

func addChild(_ value: T) -> TrieNode? {
guard children[value] == nil else {
return nil
}
children[value] = TrieNode(value, self)
return children[value]
}
}

class Trie {
typealias Node = TrieNode<Character>
private var root: Node

/** Initialize your data structure here. */
init() {
root = Node()
}

/** Inserts a word into the trie. */
func insert(_ word: String) {
guard !word.isEmpty else { return }

var currentNode: Node? = root
var count = 0
for char in word.lowercased() {
if let child = currentNode?.children[char] {
currentNode = child
}else {
}
count += 1
if count == word.count {
currentNode?.isTerminate = true
}
}
}

/** Returns if the word is in the trie. */
func search(_ word: String) -> Bool {
guard !word.isEmpty else { return false }
var currentNode: Node? = root

for char in word.lowercased() {
currentNode = currentNode?.children[char]
if currentNode == nil { break }
}

return currentNode != nil && currentNode!.isTerminate
}

/** Returns if there is any word in the trie that starts with the given prefix. */
func startsWith(_ prefix: String) -> Bool {
guard !prefix.isEmpty else { return false }
var currentNode: Node? = root

for char in prefix.lowercased() {
currentNode = currentNode?.children[char]
if currentNode == nil { break }
}

return currentNode != nil
}
}
``````