[Array] Optimize and add a solution to Shortest Word Distance i and iii

String/ShortestWordDistance.swift → Array/ShortestWordDistance.swift

@@ -7,21 +7,18 @@
 class ShortestWordDistance {
     func shortestDistance(_ words: [String], _ word1: String, _ word2: String) -> Int {
         var distance = Int.max
-        var firstIndex: Int?
-        var secondIndex: Int?
+        var firstIndex = -1, secondIndex = -1
 
         for (i, word) in words.enumerated() {
             if word == word1 {
-                if let sec = secondIndex {
-                    distance = min(distance, i - sec)
-                }
                 firstIndex = i
-            } else if word == word2 {
-                if let fir = firstIndex {
-                    distance = min(distance, i - fir)
-                }
+            } 
+            if word == word2 {
                 secondIndex = i
             }
+            if firstIndex != -1 && secondIndex != -1 {
+                distance = min(distance, abs(firstIndex - secondIndex))
+            }
         }
 
         return distance

Array/ShortestWordDistanceIII.swift

@@ -0,0 +1,34 @@
+/**
+ * Question Link: https://leetcode.com/problems/shortest-word-distance-iii/
+ * Primary idea: Iterate and update index and distance when encounter word1 or word2, use
+ *			     a temp variable to memorize the previous postion if word1 == word2
+ *
+ * Time Complexity: O(n), Space Complexity: O(1)
+ */
+
+ class ShortestWordDistanceIII {
+    func shortestWordDistance(_ words: [String], _ word1: String, _ word2: String) -> Int {
+        var idx1 = -1, idx2 = -1, res = Int.max
+
+        for (i, word) in words.enumerated() {
+            var prev = idx1
+
+            if word == word1 {
+                idx1 = i
+            }
+            if word == word2 {
+                idx2 = i
+            }
+
+            if idx1 != -1 && idx2 != -1 {
+                if word1 == word2 && prev != -1 && prev != idx1 {
+                    res = min(res, idx1 - prev)
+                } else if idx1 != idx2 {
+                    res = min(res, abs(idx1 - idx2))
+                }
+            }
+        }
+
+        return res
+    }
+}