Merge pull request #383 from abhisood/abhisood-wordsInSubtrie-partialWord

Trie - Support prefix matching

Author: kelvinlauKL

Trie/Trie/Trie/Trie.swift

@@ -135,22 +135,38 @@ extension Trie {
     }
     return currentNode.isTerminating
   }
-  
+
+
+  /// Attempts to walk to the last node of a word.  The
+  /// search will fail if the word is not present. Doesn't
+  /// check if the node is terminating
+  ///
+  /// - Parameter word: the word in question
+  /// - Returns: the node where the search ended, nil if the
+  /// search failed.
+  private func findLastNodeOf(word: String) -> Node? {
+      var currentNode = root
+      for character in word.lowercased().characters {
+          guard let childNode = currentNode.children[character] else {
+              return nil
+          }
+          currentNode = childNode
+      }
+      return currentNode
+  }
+
   /// Attempts to walk to the terminating node of a word.  The
   /// search will fail if the word is not present.
   ///
   /// - Parameter word: the word in question
   /// - Returns: the node where the search ended, nil if the
   /// search failed.
   private func findTerminalNodeOf(word: String) -> Node? {
-    var currentNode = root
-    for character in word.lowercased().characters {
-      guard let childNode = currentNode.children[character] else {
-        return nil
-      }
-      currentNode = childNode
+    if let lastNode = findLastNodeOf(word: word) {
+        return lastNode.isTerminating ? lastNode : nil
     }
-    return currentNode.isTerminating ? currentNode : nil
+    return nil
+
   }
 
   /// Deletes a word from the trie by starting with the last letter
@@ -201,7 +217,7 @@ extension Trie {
   ///   - rootNode: the root node of the subtrie
   ///   - partialWord: the letters collected by traversing to this node
   /// - Returns: the words in the subtrie
-  func wordsInSubtrie(rootNode: Node, partialWord: String) -> [String] {
+  fileprivate func wordsInSubtrie(rootNode: Node, partialWord: String) -> [String] {
     var subtrieWords = [String]()
     var previousLetters = partialWord
     if let value = rootNode.value {
@@ -216,4 +232,25 @@ extension Trie {
     }
     return subtrieWords
   }
+
+  /// Returns an array of words in a subtrie of the trie that start
+  /// with given prefix
+  ///
+  /// - Parameters:
+  ///   - prefix: the letters for word prefix
+  /// - Returns: the words in the subtrie that start with prefix
+  func findWordsWithPrefix(prefix: String) -> [String] {
+      var words = [String]()
+      let prefixLowerCased = prefix.lowercased()
+      if let lastNode = findLastNodeOf(word: prefixLowerCased) {
+          if lastNode.isTerminating {
+              words.append(prefixLowerCased)
+          }
+          for childNode in lastNode.children.values {
+              let childWords = wordsInSubtrie(rootNode: childNode, partialWord: prefixLowerCased)
+              words += childWords
+          }
+      }
+      return words
+  }
 }

Trie/Trie/TrieTests/TrieTests.swift

@@ -168,4 +168,28 @@ class TrieTests: XCTestCase {
         XCTAssertEqual(trieCopy.count, trie.count)
 
     }
+
+    func testFindWordsWithPrefix() {
+        let trie = Trie()
+        trie.insert(word: "test")
+        trie.insert(word: "another")
+        trie.insert(word: "exam")
+        let wordsAll = trie.findWordsWithPrefix(prefix: "")
+        XCTAssertEqual(wordsAll.sorted(), ["another", "exam", "test"]);
+        let words = trie.findWordsWithPrefix(prefix: "ex")
+        XCTAssertEqual(words, ["exam"]);
+        trie.insert(word: "examination")
+        let words2 = trie.findWordsWithPrefix(prefix: "exam")
+        XCTAssertEqual(words2, ["exam", "examination"]);
+        let noWords = trie.findWordsWithPrefix(prefix: "tee")
+        XCTAssertEqual(noWords, []);
+        let unicodeWord = "😬😎"
+        trie.insert(word: unicodeWord)
+        let wordsUnicode = trie.findWordsWithPrefix(prefix: "😬")
+        XCTAssertEqual(wordsUnicode, [unicodeWord]);
+        trie.insert(word: "Team")
+        let wordsUpperCase = trie.findWordsWithPrefix(prefix: "Te")
+        XCTAssertEqual(wordsUpperCase.sorted(), ["team", "test"]);
+
+    }
 }