@@ -135,22 +135,38 @@ extension Trie {
135135 }
136136 return currentNode. isTerminating
137137 }
138-
138+
139+
140+ /// Attempts to walk to the last node of a word. The
141+ /// search will fail if the word is not present. Doesn't
142+ /// check if the node is terminating
143+ ///
144+ /// - Parameter word: the word in question
145+ /// - Returns: the node where the search ended, nil if the
146+ /// search failed.
147+ private func findLastNodeOf( word: String ) -> Node ? {
148+ var currentNode = root
149+ for character in word. lowercased ( ) . characters {
150+ guard let childNode = currentNode. children [ character] else {
151+ return nil
152+ }
153+ currentNode = childNode
154+ }
155+ return currentNode
156+ }
157+
139158 /// Attempts to walk to the terminating node of a word. The
140159 /// search will fail if the word is not present.
141160 ///
142161 /// - Parameter word: the word in question
143162 /// - Returns: the node where the search ended, nil if the
144163 /// search failed.
145164 private func findTerminalNodeOf( word: String ) -> Node ? {
146- var currentNode = root
147- for character in word. lowercased ( ) . characters {
148- guard let childNode = currentNode. children [ character] else {
149- return nil
150- }
151- currentNode = childNode
165+ if let lastNode = findLastNodeOf ( word: word) {
166+ return lastNode. isTerminating ? lastNode : nil
152167 }
153- return currentNode. isTerminating ? currentNode : nil
168+ return nil
169+
154170 }
155171
156172 /// Deletes a word from the trie by starting with the last letter
@@ -201,7 +217,7 @@ extension Trie {
201217 /// - rootNode: the root node of the subtrie
202218 /// - partialWord: the letters collected by traversing to this node
203219 /// - Returns: the words in the subtrie
204- func wordsInSubtrie( rootNode: Node , partialWord: String ) -> [ String ] {
220+ fileprivate func wordsInSubtrie( rootNode: Node , partialWord: String ) -> [ String ] {
205221 var subtrieWords = [ String] ( )
206222 var previousLetters = partialWord
207223 if let value = rootNode. value {
@@ -216,4 +232,24 @@ extension Trie {
216232 }
217233 return subtrieWords
218234 }
235+
236+ /// Returns an array of words in a subtrie of the trie that start
237+ /// with given prefix
238+ ///
239+ /// - Parameters:
240+ /// - prefix: the letters for word prefix
241+ /// - Returns: the words in the subtrie that start with prefix
242+ func findWordsWithPrefix( prefix: String ) -> [ String ] {
243+ var words = [ String] ( )
244+ if let lastNode = findLastNodeOf ( word: prefix) {
245+ if lastNode. isTerminating {
246+ words. append ( prefix)
247+ }
248+ for childNode in lastNode. children. values {
249+ let childWords = wordsInSubtrie ( rootNode: childNode, partialWord: prefix)
250+ words += childWords
251+ }
252+ }
253+ return words
254+ }
219255}
0 commit comments