Updates implementation.

Author: kelvinlauKL

3Sum and 4Sum/3Sum.playground/Contents.swift

@@ -1,59 +1,76 @@
-func ThreeSum(_ nums: [Int], targetSum: Int) -> [[Int]] {
-    var a = nums.sorted()
-    var ret: [[Int]] = []
-    
-    for m in 0..<a.count {
-        if m != 0 && a[m] == a[m-1] {
-            continue
-        }
-        
-        var l = m + 1
-        var r = nums.count - 1
-        
-        while l < r {
-            let sum = a[m] + a[l] + a[r]
-            
-            if sum == targetSum {
-                ret.append([a[m], a[l], a[r]])
-                l += 1
-            } else if sum < targetSum {
-                l += 1
-            } else {
-                r -= 1
-            }
-            
-            if sum == targetSum {
-                while l < r {
-                    var flag = false
-                    if l != 0 && a[l] == a[l-1] {
-                        l += 1
-                        flag = true
-                    }
-                    
-                    if (r != a.count - 1) && a[r] == a[r+1] {
-                        r -= 1
-                        flag = true
-                    }
-                    
-                    if (!flag) {
-                        break
-                    }
-                }
-            }
-        }
+extension Collection where Element: Equatable {
+  
+  /// Returns next index with unique value. Works only on sorted arrays.
+  ///
+  /// - Parameter index: The current `Int` index.
+  /// - Returns: The new `Int` index. Will return `nil` if new index happens to be the `endIndex` (out of bounds)
+  func uniqueIndex(after index: Index) -> Index? {
+    guard index < endIndex else { return nil }
+    var index = index
+    var nextIndex = self.index(after: index)
+    while nextIndex < endIndex && self[index] == self[nextIndex] {
+      formIndex(after: &index)
+      formIndex(after: &nextIndex)
     }
 
-    return ret
+    if nextIndex == endIndex {
+      return nil
+    } else {
+      return nextIndex
+    }
+  }
 }
 
-// Answer: [[-1, 0, 1], [-1, -1, 2]]
-ThreeSum([-1, 0, 1, 2, -1, -4], targetSum: 0)
+extension BidirectionalCollection where Element: Equatable {
+  
+  /// Returns next index with unique value. Works only on sorted arrays.
+  ///
+  /// - Parameter index: The current `Int` index.
+  /// - Returns: The new `Int` index. Will return `nil` if new index happens to come before the `startIndex` (out of bounds)
+  func uniqueIndex(before index: Index) -> Index? {
+    return indices[..<index].reversed().first { index -> Bool in
+      let nextIndex = self.index(after: index)
+      guard nextIndex >= startIndex && self[index] != self[nextIndex] else { return false }
+      return true
+    }
+  }
+}
 
-// Answer: [[-1, -1, 1]]
-ThreeSum([-1, 1, 1, -1], targetSum: 1)
+func threeSum<T: BidirectionalCollection>(_ collection: T, target: T.Element) -> [[T.Element]] where T.Element: BinaryInteger & Comparable {
+  let sorted = collection.sorted()
+  var ret: [[T.Element]] = []
+  
+  for l in sequence(first: sorted.startIndex, next: sorted.uniqueIndex(after:)) {
+    var m = sorted.index(after: l)
+    var r = sorted.index(before: sorted.endIndex)
+
+    while m < r {
+      let sum = sorted[l] + sorted[m] + sorted[r]
+      switch target {
+      case sum:
+        ret.append([sorted[l], sorted[m], sorted[r]])
+        guard let nextMid = sorted.uniqueIndex(after: m), let nextRight = sorted.uniqueIndex(before: r) else { break }
+        m = nextMid
+        r = nextRight
+      case ..<target:
+        guard let nextMid = sorted.uniqueIndex(after: m) else { break }
+        m = nextMid
+      case target...:
+        guard let nextRight = sorted.uniqueIndex(before: r) else { break }
+        r = nextRight
+      default: fatalError("Swift isn't smart enough to detect that this switch statement is exhausive")
+      }
+    }
+  }
+  
+  return ret
+}
+
+// Answer: [[-1, 0, 1], [-1, -1, 2]]
+threeSum([-1, 0, 1, 2, -1, -4], target: 0)
 
 // Answer: [[-1, -1, 2], [-1, 0, 1]]
-ThreeSum([-1, -1, -1, -1, 2, 1, -4, 0], targetSum: 0)
+threeSum([-1, -1, -1, -1, 2, 1, -4, 0], target: 0)
 
 // Answer: [[-1, -1, 2]]
-ThreeSum([-1, -1, -1, -1, -1, -1, 2], targetSum: 0)
+threeSum([-1, -1, -1, -1, -1, -1, 2], target: 0)