[DFS] Refactor solution to Word Squares

Author: soapyigu

DFS/WordSquares.swift

@@ -2,61 +2,64 @@
  * Question Link: https://leetcode.com/problems/word-squares/
  * Primary idea: Classic Depth-first Search, fill out row by row, choose correct word with fixed prefix, only need to care which column is used
  * 
- * Time Complexity: O(n^n), Space Complexity: O(n)
+ * Time Complexity: O(n^2), Space Complexity: O(n^2)
  *
  */
 
  class WordSquares {
     func wordSquares(_ words: [String]) -> [[String]] {
-        var paths = [[String]]()
-        
-        guard words.count > 0 else {
-            return paths
+        guard let first = words.first else {
+            return [[String]]()
         }
 
-        var prefixWords = initPrefix(words), path = [String]()
+        let prefixesToWords = buildMaps(words), rowNum = first.count
+        var paths = [[String]](), path = [String]()
 
-        dfs(&paths, &path, prefixWords, 0, words[0].count)
+        dfs(&paths, path, prefixesToWords, rowNum, 0)
 
         return paths
     }
 
-    fileprivate func dfs(_ paths: inout [[String]], _ path: inout [String], _ prefixWords: [String:[String]], _ row: Int, _ len: Int) {
-        if row == len {
-            paths.append(Array(path))
+    private func dfs(_ paths: inout [[String]], _ path: [String], _ prefixesToWords: [String: [String]], _ rowNum: Int, _ currentRow: Int) {
+        if currentRow == rowNum {
+            paths.append(path)
             return
         }
 
-        var pre = ""
-        for i in 0..<row {
-            pre = "\(pre)\(Array(path[i])[row])"
+        var prefix = ""
+        for i in 0..<currentRow {
+            prefix.append(Array(path[i])[currentRow])
         }
 
-        guard let words = prefixWords[pre] else {
+        guard let words = prefixesToWords[prefix] else {
             return
         }
 
         for word in words {
-            path.append(word)
-            dfs(&paths, &path, prefixWords, row + 1, len)
-            path.removeLast()
+            dfs(&paths, path + [word], prefixesToWords, rowNum, currentRow + 1)
         }
     }
 
-    fileprivate func initPrefix(_ words: [String]) -> [String: [String]]{
-        var prefix = [String: [String]]()
+    private func buildMaps(_ words: [String]) -> [String: [String]] {
+        var res = [String: [String]]()
 
         for word in words {
-            prefix[""] = prefix["", default:[]] + [word]
-            
-            var pre = ""
-            for c in word {
-                pre = "\(pre)\(c)"
-                
-                prefix[pre] = prefix[pre, default:[]] + [word]
+            for prefix in prefixes(word) {
+                res[prefix, default: []].append(word)
             }
         }
 
-        return prefix
+        return res
+    }
+    
+    private func prefixes(_ word: String) -> [String] {
+        var res = [""], prefix = ""
+        
+        for char in word {
+            prefix.append(char)
+            res.append(prefix)
+        }
+        
+        return res
     }
 }

README.md

@@ -293,7 +293,7 @@
 [Partition to K Equal Sum Subsets](https://leetcode.com/problems/partition-to-k-equal-sum-subsets/)| [Swift](./DFS/PartitionKEqualSumSubsets.swift)| Medium| O(k^n)| O(n)|
 [N-Queens](https://leetcode.com/problems/n-queens/)| [Swift](./DFS/NQueens.swift)| Hard| O((n^n))| O(n^2)|
 [N-Queens II](https://leetcode.com/problems/n-queens-ii/)| [Swift](./DFS/NQueensII.swift)| Hard| O((n^n))| O(n)|
-[Word Squares](https://leetcode.com/problems/word-squares/)| [Swift](./DFS/WordSquares.swift)| Hard| O((n^n))| O(n)|
+[Word Squares](https://leetcode.com/problems/word-squares/)| [Swift](./DFS/WordSquares.swift)| Hard| O((n^2))| O(n^2)|
 [Word Pattern II](https://leetcode.com/problems/word-pattern-ii/)| [Swift](./DFS/WordPatternII.swift)| Hard| O(n^n)| O(n)|
 [Sudoku Solver](https://leetcode.com/problems/sudoku-solver/)| [Swift](./DFS/SudokuSolver.swift)| Hard| O(n^4)| O(1)|
 [Remove Invalid Parentheses](https://leetcode.com/problems/remove-invalid-parentheses/)| [Swift](./DFS/RemoveInvalidParentheses.swift)| Hard| O(n^n)| O(n)|