Create 658-Find-K-Closest-Elements.py

Author: neetcode-gh

658-Find-K-Closest-Elements.py

@@ -0,0 +1,42 @@
+# Log(n) + k
+# 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
+        
+        # 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
+
+            if arr[m] < x:    l = m + 1
+            elif arr[m] > x:  r = m - 1
+            else:             break
+    
+        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]
+      
+# Log(n-k) + k
+# Elegant but very difficult to understand
+class Solution:
+    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
+        l, r = 0, len(arr) - k
+        
+        while l < r:
+            m = (l + r) // 2
+            if x - arr[m] > arr[m + k] - x:
+                l = m + 1
+            else: 
+                r = m
+        return arr[l:l+k]