Adding an improved solution for Problem: 658. Find K Closest Elements

python/0658-find-k-closest-elements.py

@@ -2,32 +2,32 @@
 # More code but also more intuitive
 class Solution:
     def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
-        l, r = 0, len(arr) - 1
+        # Step 1: Binary Search to find the closest position to x
+        left, right = 0, len(arr) - 1
+        while left < right:
+            mid = (left + right) // 2
+            if arr[mid] < x:
+                left = mid + 1
+            else:
+                right = mid
 
-        # Find index of x or the closest val to x
-        val, idx = arr[0], 0
-        while l <= r:
-            m = (l + r) // 2
-            curDiff, resDiff = abs(arr[m] - x), abs(val - x)
-            if curDiff < resDiff or (curDiff == resDiff and arr[m] < val):
-                val, idx = arr[m], m
+        # Step 2: Initialize two pointers around the closest element
+        left, right = left - 1, left  # left points to smaller, right points to larger
 
-            if arr[m] < x:
-                l = m + 1
-            elif arr[m] > x:
-                r = m - 1
+        # Step 3: Expand around the closest position to find k elements
+        while k > 0:
+            if left < 0:  # If left pointer is out of bounds, move right
+                right += 1
+            elif right >= len(arr):  # If right pointer is out of bounds, move left
+                left -= 1
+            elif abs(arr[left] - x) <= abs(arr[right] - x):  # Pick the closer element
+                left -= 1
             else:
-                break
+                right += 1
+            k -= 1
 
-        l = r = idx
-        for i in range(k - 1):
-            if l == 0:
-                r += 1
-            elif r == len(arr) - 1 or x - arr[l - 1] <= arr[r + 1] - x:
-                l -= 1
-            else:
-                r += 1
-        return arr[l : r + 1]
+        # Step 4: Return the sorted subarray
+        return arr[left + 1:right]
 
 
 # Log(n-k) + k