[DFS] Update Trie Implementation for Word Search II

dyljqq · Feb 13, 2018 · 518b7f6 · 518b7f6
1 parent c3799cf
commit 518b7f6
Showing 1 changed file with 48 additions and 56 deletions.
diff --git a/DFS/WordSearchII.swift b/DFS/WordSearchII.swift
@@ -67,76 +67,68 @@ class WordSearchII {
 
         return trie
     }
+}
 
 
-    class Trie {
-        var root: TrieNode
+class TrieNode {
+    var isEnd: Bool
+    var letter: Character
+    var children: [Character: TrieNode]
 
-        init() {
-            root = TrieNode(letter: Character("0"))
-        }
+    init(_ letter: Character) {
+        self.isEnd = false
+        self.letter = letter
+        self.children = [Character: TrieNode]()
+    }
+}
+
+
+class Trie {
+    var root: TrieNode?
+
+    init() {
+        root = TrieNode(Character("0"))
+    }
 
-        func insert(word: String) {
-            var node = root
-            var word = [Character](word.characters)
+    func insert(_ word: String) {
+        var node = root
 
-            for i in 0..<word.count {
-                let c = word[i]
-
-                if node.children[c] == nil {
-                    node.children[c] = TrieNode(letter: c)
-                }
-
-                node = node.children[c]!
+        for char in word {
+            if node?.children[char] == nil {
+                node?.children[char] = TrieNode(char)
             }
-
-            node.isEnd = true
+      
+            node = node?.children[char]
         }
+
+        node?.isEnd = true
+    }
 
-        func isWord(word: String) -> Bool {
-            var node = root
-            var word = [Character](word.characters)
+    func search(_ word: String) -> Bool {
+        var node = root
 
-            for i in 0..<word.count {
-                let c = word[i]
-
-                if node.children[c] == nil {
-                    return false
-                }
-
-                node = node.children[c]!
+        for char in word {
+            guard let childTrieNode = node?.children[char] else {
+                return false
             }
-
-            return node.isEnd
-        }
-
-        func isWordPrefix(prefix: String) -> Bool {
-            var node = root
-            var prefix = [Character](prefix.characters)
-
-            for i in 0..<prefix.count {
-                let c = prefix[i]
-
-                if node.children[c] == nil {
-                    return false
-                }
 
-                node = node.children[c]!
-            }
-
-            return true
+            node = childTrieNode
         }
+
+        return node?.isEnd
     }
-
-    class TrieNode {
-        var letter: Character
-        var isEnd: Bool
-        var children: [Character:TrieNode]
 
-        init(letter: Character) {
-            self.letter = letter
-            isEnd = false
-            children = [Character:TrieNode]()
+    func start(with word: String) -> Bool {
+        var node = root
+
+        for char in word {
+            guard let childTrieNode = node?.children[char] else {
+                return false
+            }
+
+            node = childTrieNode
         }
+
+        return true
     }
 }