Refactor code to offer better performance and readability

Author: Yi Gu

DFS/WordSearchII.swift

@@ -11,131 +11,78 @@
 class WordSearchII {
 
     func findWords(_ board: [[Character]], _ words: [String]) -> [String] {
-        var res = [String]()
-  
-        let m = board.count
-        let n = board[0].count
-
-        let trie = _convertToTrie(words)
-        var visited = [[Bool]](repeating: Array(repeating: false, count: n), count: m)
-  
-        for i in 0 ..< m {
-            for j in 0 ..< n {
-                _dfs(board, m, n, i, j, &visited, &res, trie, "")
+        let trie = Trie(words), m = board.count, n = board[0].count
+        var isVisited = Array(repeating: Array(repeating: false, count: n), count: m), res = Set<String>()
+        
+        for i in 0..<m {
+            for j in 0..<n {
+                search(board, trie, &res, i, j, &isVisited, trie.root, "", m, n)
             }
         }
-  
-        return res
+        
+        return Array(res)
     }
 
-    fileprivate func _dfs(_ board: [[Character]], _ m: Int, _ n: Int, _ i: Int, _ j: Int, _ visited: inout [[Bool]], _ res: inout  [String], _ trie: Trie, _ str: String) {
-        // beyond matrix
+    private func search(_ board: [[Character]], _ trie: Trie, _ res: inout Set<String>, _ i: Int, _ j: Int, _ isVisited: inout [[Bool]], _ currentNode: TrieNode, _ currentStr: String, _ m: Int, _ n: Int) {
         guard i >= 0 && i < m && j >= 0 && j < n else {
             return
         }
-  
-        // check visited
-        guard !visited[i][j] else {
+        
+        guard !isVisited[i][j] else {
             return
         }
-  
-        // check is word prefix
-        let str = str + "\(board[i][j])"
-        guard trie.isWordPrefix(str) else {
+        
+        guard let child = currentNode.children[board[i][j]] else {
             return
         }
-  
-        // check word exist
-        if trie.isWord(str) && !res.contains(str) {
-            res.append(str)
+        
+        isVisited[i][j] = true
+            
+        let str = currentStr + "\(board[i][j])"
+            
+        if child.isEnd {
+            res.insert(str)
         }
-  
-        // check four directions
-        visited[i][j] = true
-        _dfs(board, m, n, i + 1, j, &visited, &res, trie, str)
-        _dfs(board, m, n, i - 1, j, &visited, &res, trie, str)
-        _dfs(board, m, n, i, j + 1, &visited, &res, trie, str)
-        _dfs(board, m, n, i, j - 1, &visited, &res, trie, str)
-        visited[i][j] = false
+        
+        search(board, trie, &res, i + 1, j, &isVisited, child, str, m, n)
+        search(board, trie, &res, i - 1, j, &isVisited, child, str, m, n)
+        search(board, trie, &res, i, j + 1, &isVisited, child, str, m, n)
+        search(board, trie, &res, i, j - 1, &isVisited, child, str, m, n)
+            
+        isVisited[i][j] = false
     }
 
-    func _convertToTrie(_ words: [String]) -> Trie {
-        let trie = Trie()
-  
-        for str in words {
-            trie.insert(str)
+    class Trie {
+        var root: TrieNode
+        
+        init(_ words: [String]) {
+            root = TrieNode()
+            
+            words.forEach { insert($0) }
         }
-  
-        return trie
-    }
-}
-
-    
-class Trie {
-    var root: TrieNode
-  
-    init() {
-        root = TrieNode()
-    }
-  
-    func insert(_ word: String) {
-        var node = root
-        var word = [Character](word.characters)
-    
-        for i in 0 ..< word.count {
-            let c = word[i]
-      
-            if node.children[c] == nil {
-                node.children[c] = TrieNode()
+        
+        private func insert(_ word: String) {
+            var node = root
+            
+            for char in word {
+                if node.children[char] == nil {
+                    node.children[char] = TrieNode()
+                }
+                
+                node = node.children[char]!
             }
-      
-            node = node.children[c]!
+            
+            node.isEnd = true
         }
-    
-        node.isEnd = true
     }
-  
-    func isWord(_ word: String) -> Bool {
-        var node = root
-        var word = [Character](word.characters)
-    
-        for i in 0 ..< word.count {
-            let c = word[i]
-      
-            if node.children[c] == nil {
-                return false
-            }
-      
-            node = node.children[c]!
-        }
 
-        return node.isEnd
-    }
+    class TrieNode {
+        var isEnd: Bool
+        var children: [Character: TrieNode]
 
-    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]!
+        init() {
+            isEnd = false
+            children = [Character: TrieNode]()
         }
-    
-        return true
     }
 }
-
-class TrieNode {
-    var isEnd: Bool
-    var children: [Character:TrieNode]
-  
-    init() {
-        isEnd = false
-        children = [Character:TrieNode]()
-    }
-}

Sort/SortColors.swift

@@ -6,23 +6,25 @@
 
 class SortColors {
     func sortColors(_ nums: inout [Int]) {
-        var red = 0, blue = nums.count - 1, i = 0
+        var redIdx = 0, blueIdx = nums.count - 1, currentIdx = 0
 
-        while i <= blue {
-            if nums[i] == 0 {
-                _swap(&nums, i, red)
-                red += 1
-                i += 1
-            } else if nums[i] == 1 {
-                i += 1
-            } else {
-                _swap(&nums, i, blue)
-                blue -= 1
+        while currentIdx <= blueIdx {
+            let num = nums[currentIdx]
+            
+            if num == 0 {
+                swap(&nums, redIdx, currentIdx)
+                redIdx += 1
+            } else if num == 2 {
+                swap(&nums, currentIdx, blueIdx)
+                blueIdx -= 1
+                currentIdx -= 1
             }
+            
+            currentIdx += 1
         }
     }
 
-    fileprivate func _swap<T>(_ nums: inout [T], _ p: Int, _ q: Int) {
-        (nums[p], nums[q]) = (nums[q], nums[p])
+    private func swap(_ nums: inout [Int], _ left: Int, _ right: Int) {
+        (nums[left], nums[right]) = (nums[right], nums[left])
     }
-}
+}