컴퓨터공학/LeetCode 1000

[LeetCode] 658. Find K Closest Elements

saurus2 2022. 9. 30. 09:43

658. Find K Closest Elements

Medium

Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

  • |a - x| < |b - x|, or
  • |a - x| == |b - x| and a < b

 

Example 1:

Input: arr = [1,2,3,4,5], k = 4, x = 3
Output: [1,2,3,4]

Example 2:

Input: arr = [1,2,3,4,5], k = 4, x = -1
Output: [1,2,3,4]

 

Constraints:

  • 1 <= k <= arr.length
  • 1 <= arr.length <= 104
  • arr is sorted in ascending order.
  • -104 <= arr[i], x <= 104

문제 풀이

  • 정렬이 되어있는 배열에서 정수 x 와 가까운 숫자들을 구하자.
  • 배열의 최대 크기가 10^4 이기 때문에 어떤 알고리즘으로 풀어도 상관없지만 문제 안에 정렬된 배열을 제공하기 때문에 이진탐색 문제로 접근해야할 확률이 크다. 
  • 조건은, x 값에서 같은 차이로 가깝다면 오른쪽 보다 왼쪽이 우선이다. 
  • 정렬 함수를 사용하는 것과 이진탐색을 이용해서 풀 수 있다. 
    • 정렬을 할때 숫자와 x 값의 차이로 정렬해준다.
    • 정렬 된 배열을 k 개 만큼 꺼내고 정렬한다.
  • 이진탐색을 사용할때는, k 개만큼 정답을 뽑을 시작점을 찾는다고 생각하자
    • 왼쪽은 배열의 처음을 가르키고, 오른쪽은 전체 크기에서 k 만큼 뺀 인덱스로 설정한다.
    • 이진탐색을 구현하되, 인덱스 이동의 조건은 x 값에서 중간 포인트 값을 뺀것이 크다면 왼쪽 탐색을 해야한다.
    • 반대로 중간 포인트에 k 만큼을 더한 인덱스에서 x 값을 뺀 것이 크다면 왼쪽을 탐색해야한다. 

소스 코드

정렬 

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        sort_arr = sorted(arr, key = lambda num: abs(x - num))
        res = []
        for i in range(k):
            res.append(sort_arr[i])
        return sorted(res)

이진탐색

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        l = 0
        r = len(arr) - k
        while l < r:
            mid = l + (r - l) // 2
            if x - arr[mid] > arr[mid + k] - x:
                l = mid + 1
            else:
                r = mid
        return arr[l : l + k]