Refactored implementation of Trie. Leaving the old implementation in a folder for future reference.

Author: Kelvin Lau

Trie/trie.swift renamed to Trie/Old Implementation/trie.swift

@@ -0,0 +1,16 @@
+let trie = Trie()
+
+trie.insert(word: "apple")
+trie.insert(word: "ap")
+trie.insert(word: "a")
+
+trie.contains(word: "apple")
+trie.contains(word: "ap")
+trie.contains(word: "a")
+
+trie.remove(word: "apple")
+trie.contains(word: "a")
+trie.contains(word: "apple")
+
+trie.insert(word: "apple")
+trie.contains(word: "apple")

Trie/Trie.playground/Contents.swift

@@ -1,64 +1,21 @@
-/// A Trie (prefix tree)
-///
-/// Some of the functionality of the trie makes use of the Queue implementation for this project
-///
-/// Every node in the Trie stores a bit of information pertaining to what it references:
-///   -Character (letter of an inserted word)
-///   -Parent (Letter that comes before the current letter in some word)
-///   -Children (Words that have more letters than available in the prefix)
-///   -isAWord (Does the current letter mark the end of a known inserted word?)
-///   -visited (Mainly for the findPrefix() function)
-public class Node {
-  public var character: String
-  public var parent: Node?
-  public var children: [String: Node]
-  public var isAWord: Bool
-  public var visited: Bool // only for findPrefix
+public final class TrieNode<T: Hashable> {
+  public var value: T?
+  public weak var parent: TrieNode?
+  public var children: [T: TrieNode] = [:]
+  public var isTerminating = false
 
-  init(character: String, parent: Node?) {
-    self.character = character
-    self.children = [:]
-    self.isAWord = false
-    self.parent = parent
-    self.visited = false
-  }
-  
-  /// Returns `true` if the node is a leaf node, false otherwise.
-  var isLeaf: Bool {
-    return children.count == 0
-  }
-  
-  /// Changes the parent of the current node to the passed in node.
-  ///
-  /// - parameter node: A `Node` object.
-  func setParent(node: Node) {
-    parent = node
-  }
+  init() {}
 
-  /// Returns the child node that holds the specific passed letter.
-  ///
-  /// - parameter character: A `String`
-  ///
-  /// - returns: The `Node` object that contains the `character`.
-  func getChildAt(character: String) -> Node {
-    return children[character]!
-  }
-  
-  /// Returns whether or not the current node marks the end of a valid word.
-  var isValidWord: Bool {
-    return isAWord
-  }
-  
-  
-  /// Returns whether or not the current node is the root of the trie.
-  var isRoot: Bool {
-    return character == ""
+  init(value: T, parent: TrieNode? = nil) {
+    self.value = value
+    self.parent = parent
   }
 }
 
-// MARK: - CustomStringConvertible
-extension Node: CustomStringConvertible {
-  public var description: String {
-    return ""
+// MARK: - Insertion
+public extension TrieNode {
+  func add(child: T) {
+    guard children[child] == nil else { return }
+    children[child] = TrieNode(value: child, parent: self)
   }
 }

Trie/Trie.playground/Sources/Node.swift

@@ -2,211 +2,82 @@
 ///   -root (the root of the trie)
 ///   -wordList (the words that currently exist in the trie)
 ///   -wordCount (the number of words in the trie)
-public class Trie {
-  private var root = Node(character: "", parent: nil)
-  private(set) var wordList: [String] = []
-  private(set) var wordCount = 0
+public final class Trie {
+  typealias Node = TrieNode<Character>
+  fileprivate let root: Node
 
-  init(words: Set<String>) {
-    words.forEach { insert(word: $0) }
+  public init() {
+    root = TrieNode<Character>()
   }
-  
-  /// Merge two `Trie` objects into one and returns the merged `Trie`
-  ///
-  /// - parameter other: Another `Trie`
-  ///
-  /// - returns: The newly unioned `Trie`.
-  func merge(other: Trie) -> Trie {
-    let newWordList = Set(wordList + other.wordList)
-    return Trie(words: newWordList)
-  }
-  
-  /// Looks for a specific key and returns a tuple that has a reference to the node (if found) and true/false depending on if it was found.
-  ///
-  /// - parameter key: A `String` that you would like to find.
-  ///
-  /// - returns: A tuple containing the an optional `Node`. If the key is found, it will return a non-nil `Node`. The second part of the tuple contains a bool indicating whether it was found or not.
-  func find(key: String) -> (node: Node?, found: Bool) {
+}
+
+// MARK: - Basic Methods
+public extension Trie {
+  func insert(word: String) {
+    guard !word.isEmpty else { return }
     var currentNode = root
 
-    for character in key.characters {
-      if currentNode.children["\(character)"] == nil {
-        return (nil, false)
-      }
-      currentNode = currentNode.children["\(character)"]!
-    }
-    return (currentNode, currentNode.isValidWord)
-  }
-  
-  /// Returns true if the `Trie` is empty, false otherwise.
-  var isEmpty: Bool {
-    return wordCount == 0
-  }
-  
-  /// Checks if the a word is currently in the `Trie`.
-  ///
-  /// - parameter word: A `String` you want to check.
-  ///
-  /// - returns: Returns `true` if the word exists in the `Trie`. False otherwise.
-  func contains(word: String) -> Bool {
-    return find(key: word.lowercased()).found
-  }
-  
-  func isPrefix(_ prefix: String) -> (node: Node?, found: Bool) {
-    let prefix = prefix.lowercased()
-    var currentNode = root
+    var characters = Array(word.lowercased().characters)
+    var currentIndex = 0
 
-    for character in prefix.characters {
-      if currentNode.children["\(character)"] == nil {
-        return (nil, false)
+    while currentIndex < characters.count {
+      let character = characters[currentIndex]
+      if let child = currentNode.children[character] {
+        currentNode = child
+      } else {
+        currentNode.add(child: character)
+        currentNode = currentNode.children[character]!
       }
 
-      currentNode = currentNode.children["\(character)"]!
-    }
-    
-    if currentNode.children.count > 0 {
-      return (currentNode, true)
+      currentIndex += 1
+      if currentIndex == characters.count {
+        currentNode.isTerminating = true
+      }
     }
-    
-    return (nil, false)
   }
 
-  /// Inserts a word into the trie. Returns a tuple containing the word attempted to tbe added, and true/false depending on whether or not the insertion was successful.
-  ///
-  /// - parameter word: A `String`.
-  ///
-  /// - returns: A tuple containing the word attempted to be added, and whether it was successful.
-  @discardableResult func insert(word: String) -> (word: String, inserted: Bool) {
-    let word = word.lowercased()
-    guard !contains(word: word), !word.isEmpty else { return (word, false) }
-    
+  func contains(word: String) -> Bool {
+    guard !word.isEmpty else { return false }
     var currentNode = root
-    var length = word.characters.count
 
-    var index = 0
-    var character = word.characters.first!
+    var characters = Array(word.lowercased().characters)
+    var currentIndex = 0
 
-    while let child = currentNode.children["\(character)"] {
+    while currentIndex < characters.count, let child = currentNode.children[characters[currentIndex]] {
       currentNode = child
-      length -= 1
-      index += 1
-      
-      if length == 0 {
-        currentNode.isAWord = true
-        wordList.append(word)
-        wordCount += 1
-        return (word, true)
-      }
-      
-      character = Array(word.characters)[index]
-    }
-    
-    let remainingCharacters = String(word.characters.suffix(length))
-    for character in remainingCharacters.characters {
-      currentNode.children["\(character)"] = Node(character: "\(character)", parent: currentNode)
-      currentNode = currentNode.children["\(character)"]!
+      currentIndex += 1
     }
 
-    currentNode.isAWord = true
-    wordList.append(word)
-    wordCount += 1
-    return (word, true)
-  }
-  
-  /// Attempts to insert all words from input array. Returns a tuple containing the input array and true if some of the words were successfully added, false if none were added.
-  ///
-  /// - parameter words: An array of `String` objects.
-  ///
-  /// - returns: A tuple stating whether all the words were inserted.
-  func insert(words: [String]) -> (wordList: [String], inserted: Bool) {
-    var successful: Bool = false
-    for word in wordList {
-      successful = insert(word: word).inserted || successful
+    if currentIndex == characters.count && currentNode.isTerminating {
+      return true
+    } else {
+      return false
     }
-    
-    return (wordList, successful)
   }
 
-  /// Removes the specified key from the `Trie` if it exists, returns tuple containing the word attempted to be removed and true/false if the removal was successful.
-  ///
-  /// - parameter word: A `String` to be removed.
-  ///
-  /// - returns: A tuple containing the word to be removed, and a `Bool` indicating whether removal was successful or not.
-  @discardableResult func remove(word: String) -> (word: String, removed: Bool) {
-    let word = word.lowercased()
-    
-    guard contains(word: word) else { return (word, false) }
+  func remove(word: String) {
+    guard !word.isEmpty else { return }
     var currentNode = root
 
-    for character in word.characters {
-      currentNode = currentNode.getChildAt(character: "\(character)")
+    var characters = Array(word.lowercased().characters)
+    var currentIndex = 0
+    
+    while currentIndex < characters.count {
+      let character = characters[currentIndex]
+      guard let child = currentNode.children[character] else { return }
+      currentNode = child
+      currentIndex += 1
     }
 
     if currentNode.children.count > 0 {
-      currentNode.isAWord = false
+      currentNode.isTerminating = false
     } else {
-      var character = currentNode.character
-      while currentNode.children.count == 0 && !currentNode.isRoot {
-        currentNode = currentNode.parent!
-        currentNode.children[character]!.parent = nil
-        character = currentNode.character
-      }
-    }
-    
-    wordCount -= 1
-    
-    var index = 0
-    for item in wordList {
-      if item == word {
-        wordList.remove(at: index)
-      }
-      index += 1
-    }
-    
-    return (word, true)
-  }
-  
-  func find(prefix: String) -> [String] {
-    guard isPrefix(prefix).found else { return [] }
-    var queue = Queue<Node>()
-    let node = isPrefix(prefix).node!
-    
-    var wordsWithPrefix: [String] = []
-    let word = prefix
-    var tmp = ""
-    queue.enqueue(element: node)
-    
-    while let current = queue.dequeue() {
-      for (_, child) in current.children {
-        if !child.visited {
-          queue.enqueue(element: child)
-          child.visited = true
-          if child.isValidWord {
-            var currentNode = child
-            while currentNode !== node {
-              tmp += currentNode.character
-              currentNode = currentNode.parent!
-            }
-            tmp = "\(tmp.characters.reversed())"
-            wordsWithPrefix.append(word + tmp)
-            tmp = ""
-          }
-        }
+      var character = currentNode.value
+      while currentNode.children.count == 0, let parent = currentNode.parent, !parent.isTerminating {
+        currentNode = parent
+        currentNode.children[character!] = nil
+        character = currentNode.value
       }
     }
-    
-    return wordsWithPrefix
-  }
-  
-  func removeAll() {
-    for word in wordList {
-      remove(word: word)
-    }
-  }
-}
-
-extension Trie: CustomStringConvertible {
-  public var description: String {
-    return ""
   }
 }