[BOJ] Python 백준 K번째 수 (11004) [feat. heap풀이, sort 풀이]

sort를 이용한 풀이 (메모리 : 622180KB, 실행시간 : 3452ms)

import sys
input = sys.stdin.readline

N, K = map(int,input().split())
A = list(map(int,input().split()))

A.sort()

print(A[K-1])

위 코드는 단순히 파이썬의 timsort를 사용해 정렬 후, K번째의 인덱스를 가져왔다.

시간복잡도

• 파이썬의 내장 sort() 함수의 시간복잡도는 평균적으로나 최악의 경우로나 O(NlogN)이다.
• 정렬 후 원소를 찾는 것은 O(1)이므로
• 따라서, 전체 시간복잡도는 정렬에 의해 결정되어 O(NlogN)이다.

 

heap을 이용한 풀이 (메모리 : 626184KB, 실행시간 : 5164ms)

import sys
import heapq
input = sys.stdin.readline

N, K = map(int,input().split())

A = list(map(int,input().split()))

hq = []
for i in A:

    if len(hq) < K:
        heapq.heappush(hq, -i)
    elif i < -hq[0]:
        heapq.heappop(hq)
        heapq.heappush(hq, -i)

print(-hq[0])

위 코드는 heap을 사용해 N개의 수들을 hq라는 배열을 최대 힙을 이용해 넣었으며, K개 갯수 이전에는 hq에 계속 넣어주고, K개 이상이 되면, 힙에서 가장 큰 수와 i 중 더 작은 수를 힙에 넣어준다.

그러면 hq[0]에 K번째 큰 수가 들어가 있게 된다.

시간복잡도

• 리스트 A의 모든 원소 N개를 순회해야 하므로 O(N)
• 루프 안에서 heappush 또는 heappop연산은 힙의 크기 K에 대해 O(logK)의 시간이 걸린다.
• 따라서, 총 시간 복잡도는 O(NlogK)가 된다.

 

그렇다면, 왜 heap 코드가 더 느릴까?

이론적으로 K가 N보다 훨씬 작을 때, O(NlogK)가 O(NlogN)보다 빠르다.

• K가 N보다 큰 경우가 존재할 수 있기 때문이다.
• A.sort는 상수 요인이 매우 작은, 고도로 최적화된 구현체이다. 반면, 파이썬 루프를 도는 heapq코드는 상수 요인이 상대적으로 크다
• 파이썬의 기본 정렬 알고리즘은 TimSort는 이미 데이터가 일부 정렬되어 있는 경우 등 실제 데이터에서 나타나는 특징을 잘 활용하여 O(NlogN)보다 훨씬 빠르게 거의 O(N)에 가깝게 동작한다.