[python] K번 째 수 찾기

import heapq

def findKth(myInput, k) :
    '''
    매 순간마다 k번째로 작은 원소를 리스트로 반환합니다.
    '''
    result = []
    l = [] #max_heap
    r = [] #min_heap
    
    for i in myInput:
    	# max_heap이 비어있거나, max_heap의 루트값보다 i값이 작거나 같다면, -i 삽입
        if not l or -l[0] >= i:
            heapq.heappush(l, -i)
        # 그 외의 경우, min_heap에 i 삽입
        else:
            heapq.heappush(r, i)
        
        # min_heap의 루트노드(가장 큰 값)을 K번째 작은 수로 만들기 위하여 min_heap의 크기를 k개로 유지
        if len(l) > k:
            heapq.heappush(r, -heapq.heappop(l))
        elif len(l) < k and r:
            heapq.heappush(l, -heapq.heappop(r))
        
        result.append(-l[0] if len(l) >= k else -1) 

    return result

어떤 배열이 주어졌을 때 K번째 작은 값을 구하는 코드이다. min_heap, max_heap 총 두개의 힙 자료구조를 사용하여 O(nlogn)의 시간복잡도(n 개의 원소, 삽입시 O(logn)의 시간 복잡도)로 구할 수 있다.

어떤 배열이 주어졌을 때 단순히 quick sort나 merge sort로 O(nlogn)의 시간 복잡도로 정렬 후 인덱스로 접근하는 방법도 있지만 배열에 새로운 값이 추가되는 상황에서 계속해서 K번째 값을 추적해야 하는 경우, 정렬 방법은 원소가 추가될 때마다 O(nlogn)의 시간복잡도를 갖는 반면, 힙을 사용하면 O(logn)의 시간복잡도로 처리가 가능하다는 장점이 있다.

 

위의 로직을 응용하여 배열의 중간값을 구할 수도 있다. 

참고자료 - https://unfunhy.tistory.com/44

백준 2696 : 중앙값 구하기