오늘의 한 마디
중위값을 이렇게 구하는 게 참신하네

• Baekjoon 문제 링크
• 현재 백준 문제집: 단기간 성장을 풀고 있습니다.

발상

처음에는 넣기만 하면 자동으로 정렬된 상태를 만들어주는 자료구조가 있지 않을까 라는 생각을 했다.

저런 자료구조가 있다면, 숫자를 1, 2, 3, 4, 5개 부를 때마다 중위값의 인덱스는 1, 1, 2, 2, 3, ... 이런 규칙으로 증가하기 때문에 시간복잡도를 줄일수 있겠다고 생각했다.

하지만 그런 자료구조는 파이썬에 없었다.

생각나는 건 heapq 밖에 없었기 때문에...

LinkedList로 만들면 어떨까 생각해봤는데, LinkedList는 인덱스로 접근하는 방식이 아니라 매번 head부터 Runner 기법을 사용하여 O(n)의 탐색을 해야 하기 때문에 시간복잡도 상 얻는 이익이 없다.

해설

포인트는 중위값보다 작거나 같은 힙과 중위값보다 큰 힙, 총 두개의 힙을 만드는 것이다.

중위값보다 작거나 같은 힙(leftq)은 최대힙으로, 중위값보다 큰 힙(rightq)은 최소힙으로 만든다. leftq의 최댓값은 뱀의 머리, rightq의 최솟값은 용의 꼬리에 비유할 수 있다.

작동 과정은 다음과 같다.

1. leftq와 rightq의 길이가 같다면 leftq에 넣고,
   그렇지 않다면(leftq에 1개 더 많이 들어있는 경우밖에 없음) rightq에 넣는다.

2. 하지만 새로 넣은 요소가 중위값보다 큰데 leftq에 들어갈 수도 있고, 중위값보다 작은데 rightq에 들어갈 수도 있기 때문에 leftq[0]과 rightq[0]을 비교하여 rightq[0]이 더 크다면 둘을 바꿔줘야 한다.

하지만, 무턱대고 swap했다가 틀렸습니다를 받았다.

import sys
import heapq

N = int(sys.stdin.readline().rstrip())
leftq = [] # 최대 힙
rightq = [] # 최소 힙

for _ in range(N):
    n = int(sys.stdin.readline().rstrip())
    if len(leftq) == len(rightq):
        heapq.heappush(leftq, -n)
    else:
        heapq.heappush(rightq, n)

    if leftq and rightq and -leftq[0] > rightq[0]:
        leftq[0], rightq[0] = -rightq[0], -leftq[0] # 이 부분!
    
    print(-leftq[0])

4
2
3
4
1

위의 테스트 케이스에 대해서,

2 -> [-2] [] / 답: 2
3 -> [-2] [3] / 답: 2
4 -> [-3, -2] [4] / 답: 3
1 -> [-1, -2] [3, 4] / 답: 1 / 실제 중위값은 2

마지막 1을 넣었을 때, 3과 1이 swap되었으나, 틀린 답을 출력한다.

그 이유는, 이미 알았겠지만 leftq가 [-2, -1]로 heapify되어야 하는데
단순히 [0]번 요소를 바꾸기만 해서 그렇게 된 거다.

    if leftq and rightq and -leftq[0] > rightq[0]:
        leftq[0], rightq[0] = -rightq[0], -leftq[0]

위 코드를 아래 코드로 바꾸어야 한다.

    if leftq and rightq and -leftq[0] > rightq[0]:
        heapq.heappush(leftq, -heapq.heappop(rightq))
        heapq.heappush(rightq, -heapq.heappop(leftq))

풀이

import sys
import heapq

N = int(sys.stdin.readline().rstrip())
leftq = [] # 최대 힙
rightq = [] # 최소 힙

for _ in range(N):
    n = int(sys.stdin.readline().rstrip())
    if len(leftq) == len(rightq):
        heapq.heappush(leftq, -n)
    else:
        heapq.heappush(rightq, n)

    if leftq and rightq and -leftq[0] > rightq[0]:
        heapq.heappush(leftq, -heapq.heappop(rightq))
        heapq.heappush(rightq, -heapq.heappop(leftq))
    
    print(-leftq[0])