leetcode : Find K Closest Elements

https://leetcode.com/problems/find-k-closest-elements/

Find K Closest Elements - LeetCode

 

은근히 까다로운 문제다.

이진탐색 직접 구현해서 하기는 진짜진짜 까다롭고..

bisect 함수를 쓸거면 또 off by one error를 신경써야 하고..

그냥 이진탐색 문제는 인덱스 1차이 생각하는게 진짜 귀찮다.

컴파일 없이 풀으라 그러면 못풀거같은 문제.

 

intuition은 간단하다.

bisect_left로 찾은 x와 가장 가까운 요소의 인덱스를 right pointer로 초기화하고

left pointer를 right pointer-1로 초기화한다.

 

bisect_left로 찾은 요소의 인덱스가 right pointer가 되어야 하는 이유는

[1, 3, 7, 8]

2

4

위와 같은 경우에서

bisect_left(arr, 4)의 결과값은 4가 삽입될 수 있는 위치인 2이기 때문이다.

2를 left로 둔다면 3이 right가 된다. 우린 인덱스1에 있는 요소부터 추가해나가야 되는데 left가 2, right가 3이 되어버리면.. 안된다!

 

left와 right를 비교해 더 가까운 요소를 추가하고 포인터를 변경해나가는 작업을 계속한다.

결과값은 정렬되어 있어야 하므로 heap을 이용해 저장해뒀다가 나중에 쫙 뽑아준다.

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        def select(l, r):
            if l == -1: return r, -1, r+1
            if r == len(arr): return l, l-1, len(arr)
            
            diff_l = abs(arr[l] - x)
            diff_r = abs(arr[r] - x)
            if diff_l <= diff_r:
                return l, l-1, r
            else:
                return r, l, r+1
        
        pq = []

        right = bisect_left(arr, x)
        left = right - 1
        
        for _ in range(k):
            cur, left, right = select(left,right)
            heappush(pq, arr[cur])

        answer = []
        while pq:
            answer.append(heappop(pq))
            
        return answer