LeetCode/Grind169

658. Find K Closest Elements ★★★★★

hyunkookim 2025. 5. 23. 06:36

658. Find K Closest Elements

 

정렬된 배열에서 x에 가장 가까운 원소 k개를 찾아 정렬된 상태로 반환하는 문제입니다.


✅ 핵심 아이디어

  • 배열이 정렬되어 있으므로 이분 탐색(Binary Search) 를 이용해서 빠르게 풀 수 있습니다.
  • arr[i] ~ arr[i+k-1] 범위가 정답 후보가 되도록 슬라이딩 윈도우를 이동합니다.
  • 거리 차가 작거나, 같다면 더 작은 숫자를 우선 선택합니다.

✅ 방법 1: 이분 탐색 기반 슬라이딩 윈도우 (O(log(n - k) + k))

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        # 이분 탐색으로 윈도우 시작점(start)을 찾음
        left, right = 0, len(arr) - k

        while left < right:
            mid = (left + right) // 2
            # x에 더 가까운 쪽으로 윈도우를 이동
            if abs(arr[mid] - x) > abs(arr[mid + k] - x):
                left = mid + 1
            else:
                right = mid

        # 윈도우 범위의 k개 원소 반환
        return arr[left:left + k]

 


if abs(arr[mid] - x) > abs(arr[mid + k] - x):

슬라이딩 윈도우의 위치를 조정하기 위한 핵심 로직입니다.


🔍 맥락 설명

우리는 arr[mid] ~ arr[mid + k - 1] 이라는 길이 k짜리 윈도우가 가장 x에 가까운 원소들만 포함하도록 mid를 찾고 있는 중입니다.


💡 의미 해석

이 조건은 아래와 같은 의미를 가집니다:

  • arr[mid]: 현재 윈도우의 왼쪽 끝 값
  • arr[mid + k]: 현재 윈도우의 오른쪽 바로 다음 값 (즉, 윈도우를 오른쪽으로 한 칸 이동했을 때 생기는 새로운 마지막 값)

➤ 비교 의미:

  • abs(arr[mid] - x) > abs(arr[mid + k] - x) 이면:
    → 오른쪽에 있는 값이 x에 더 가까우므로
    윈도우를 오른쪽으로 한 칸 이동하는 것이 좋다.

그래서 left = mid + 1을 합니다.


📌 직관 예시

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

step 1:
mid = 1 → window = [2, 3, 4]
check: |2 - 4| = 2 vs |5 - 4| = 1

→ 오른쪽 끝에 있는 5가 더 가까움 → 윈도우 오른쪽으로 이동
→ left = mid + 1

즉, **"왼쪽보다 오른쪽이 더 가깝다"**면 윈도우를 오른쪽으로 옮기는 조건입니다.
→ 가장 x에 가까운 연속 k개를 찾기 위한 기준이에요.

📌 예시 설명

예제 1:

arr = [1, 2, 3, 4, 5], k = 4, x = 3 → 윈도우 [1,2,3,4] 또는 [2,3,4,5] 중 [1,2,3,4] 선택됨

예제 2:

arr = [1,1,2,3,4,5], k = 4, x = -1 → [-1]에 가까운 [1,1,2,3] 선택됨

🧠 시간복잡도

  • 이분 탐색: O(log(n - k))
  • 슬라이싱: O(k)
  • 총합: O(log(n - k) + k)