BOJ1655 가운데를 말해요

문제

BOJ1655 가운데를 말해요
골드II | 백준 1655 | Python3 파이썬 풀이

알고리즘

중간값을 저장하기 위해 힙을 두 개 사용했다. 편의상 왼쪽, 오른쪽 힙이라고 지칭하겠다.

왼쪽 힙은 현재 중간값보다 작은 값이 들어가는 최대 힙이고 오른쪽 힙은 현재 중간값보다 큰 값이 들어가는 최소 힙이다.

먼저, 중간값의 정의부터 생각해보자. 중간값은 어떠한 오름차순으로 정렬 된 수열에서 중간에 위치하고 있는 값을 말한다. 이때, 수학적인 정의에서는 수열의 원소의 개수가 짝수개일 경우, 가운데 위치한 값이 두 개이므로, 두 값의 평균을 중간값으로 취하게 되지만, 이 문제에서는 더 작은 값을 취하게 된다.

AVL 트리의 아이디어를 차용해, 루트를 중간값으로 하고, 루트의 좌측 서브트리에는 루트값보다 작은 값이 들어가는 힙 트리가, 오른쪽에는 루트값보다 큰 값이 들어가는 힙 트리가 존재한다.

만약, 어떤 수열이 (1, 5, 6, 100)이라면 중간값인 5가 루트에 위치를 하고, 왼쪽 서브트리에는 1, 5로 구성된 트리, 오른쪽 서브트리에는 6, 100으로 구성된 트리가 있게 된다.
이때, 중요한 점은, 루트의 각 서브트리와 중간값은 별도로, 중간값은 항상 다른 변수에 저장한다.

이제 단계별로 생각해보자.

먼저, 맨 첫 값은 항상 왼쪽 힙에 들어가며, 중간 값은 그 값이 된다.

여기서 만약, 10이 들어오면 10은 중간값보다 크므로 오른쪽 힙에 들어가게 된다. 만약 -10이 들어온다면 중간값보다 작으므로 왼쪽 힙에 들어가게 된다.

10이 들어온 경우는 아무 문제가 없다. 중간값도 그대로 1이다. 하지만 -10이 들어온 경우 중간값이 바뀌어야 한다. AVL트리의 아이디어를 차용했다는 이유도, 루트 노드를 기준으로 좌, 우 서브트리의 balance가 깨졌기 때문에 rotate를 통해 balance를 잡아줄 것이다. 여기서는 두 힙의 원소의 개수 차이가 2 이상이 되면 balance가 깨졌다고 정의한다.

값을 전체적으로 오른쪽으로 shift한다고 생각하자. 여기서 취해야 할 연산은 간단하다.

if : 왼쪽 힙 원소 개수가 오른쪽 힙 원소 개수보다 2 이상 크다면:
    왼쪽 힙에서 pop한 값을 오른쪽 힙에 push
else:
    오른쪽 힙에서 pop한 값을 왼쪽 힙에 push

이제 중간 값을 찾아야 한다. 중간값은 어디에 위치하고 있는가? 무조건 왼쪽 힙의 루트가 중간값인가?

위쪽의 경우 값이 짝수개이다. 짝수개일 때는 무조건 (위의 shift를 통해서라도) 좌, 우의 balance가 맞게 된다. 원소의 개수가 같다는 것이다. 이 경우에는 무조건 좌측 힙의 루트 값(가장 큰 값)이 중간값이 된다. 왜냐, 전체 수열의 중간값 두 값은 각각 왼쪽 힙의 루트, 오른쪽 힙의 루트에 있게 된다. 왼쪽 힙은 현재 중간값보다 작은 값이며, 그 중에 가장 큰 값이 왼쪽 서브수열에서 가장 가운데 있게 되며, 오른쪽 힙도 반대로 같다.

홀수의 경우를 생각해보자. 홀수의 경우는 무조건 왼쪽 힙의 루트인가? 맨 처음 첫 값만 들어왔을 때는 그러했지만, 아래와 같은 경우는 오른쪽 힙의 루트가 중간값이다. 아까의 케이스 1, -10이 들어와 shift된 후, 중간값은 -10이 될 것이다. 이때, 20이 들어오면 20은 중간값보다 크므로 오른쪽 힙에 들어가게 되고 저런 경우가 생긴다. 저런 경우에도 shift를 통해 왼쪽으로 밀어주면 중간값은 항상 왼쪽 힙에 루트에 위치시킬 수 있지만, heap push pop 연산을 두 번 해야하므로, balance가 깨진 경우에만 해주기로 하고, 단순히 왼쪽 힙의 크기와 오른쪽 힙의 크기를 비교함으로써 중간값이 어디에 위치했는지 알 수 있다.

if : 전체 원소 수가 짝수:
    왼쪽 힙의 루트가 중간값
else:
    if : 왼쪽 힙 크기 > 오른쪽 힙 크기:
        왼쪽 힙의 루트가 중간값
    else:
        오른쪽 힙의 루트가 중간값

시간 복잡도

• for loop : $O(N)$
• heap push or pop : $O(\log N)$

루프 내 한 스텝 당 heap push or pop은 최소 1번, 최대 3번 일어나므로 시간 복잡도는 $O(N\log N)$

최대인 3번 일어나는 경우는 두 힙의 개수가 2 이상 차이가 나서 shift 하는 경우인데 이 경우는 최소 2step마다, 최대는 무한 step마다 일어난다.

코드

import sys
from heapq import heappop, heappush

input = sys.stdin.readline

N = int(input())

leftpq = []
leftnum = 0
rightpq = []
rightnum = 0

m = int(input())
print(m)

# 첫 값은 무조건 왼쪽 힙에 넣는다.
heappush(leftpq, -m)
leftnum += 1

for i in range(1, N):
    X = int(input())
    
    # 중간값보다 크다면
    if X >= m:
    	# 오른쪽 힙에 넣음
        heappush(rightpq, X)
        rightnum += 1
	
    # 중간값보다 작다면
    else:
    	# 왼쪽 힙에 넣음
        heappush(leftpq, -X)
        leftnum += 1

    # -> right shift : 왼쪽 힙이 오른쪽 힙보다 2개 더 많다면
    if leftnum - rightnum == 2:
        heappush(rightpq, -(heappop(leftpq))) # shift
        leftnum -= 1 # 개수 조정
        rightnum += 1

    # <- left shift : 오른쪽 힙이 왼쪽 힙보다 2개 더 많다면
    elif rightnum - leftnum == 2:
        heappush(leftpq, -(heappop(rightpq))) # shift
        leftnum += 1 # 개수 조정
        rightnum -= 1
	
    # 중간값 찾기
    # 총 원소 개수가 짝수개면
    if (leftnum + rightnum) % 2 == 0:
        m = -(leftpq[0]) # 중간값은 무조건 왼쪽힙의 루트
    
    # 홀수개라면
    else:
    	# 왼쪽 힙 크기가 크면
        if leftnum > rightnum:
            # 왼쪽 힙의 루트가 중간값
            m = -(leftpq[0])
        else:
            # 오른쪽 힙의 루트가 중간값
            m = rightpq[0]

    # 중간값 출력
    print(m)
   

결과