[Python] G2_1655_가운데를 말해요 ⏺

https://www.acmicpc.net/problem/1655

문제

백준이는 동생에게 "가운데를 말해요" 게임을 가르쳐주고 있다. 백준이가 정수를 하나씩 외칠때마다 동생은 지금까지 백준이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 백준이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 백준이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 백준이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에는 백준이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 백준이가 외치는 정수가 차례대로 주어진다. 정수는 -10,000보다 크거나 같고, 10,000보다 작거나 같다.

출력

한 줄에 하나씩 N줄에 걸쳐 백준이의 동생이 말해야 하는 수를 순서대로 출력한다.

예제

조건

• 시간 제한: 0.1초
• 메모리 제한: 128MB

코드

import heapq,sys
input = sys.stdin.readline

N = int(input())
left_hq,right_hq = [],[]

for _ in range(N):
    num = int(input())
    
    # 길이가 같을 경우, 최대힙(=left heap)에 값을 추가
    if len(left_hq) == len(right_hq):
        heapq.heappush(left_hq,(-num,num))
    # 길이가 다른 경우, 최소힙(=right heap)에 값을 추가
    else:
        heapq.heappush(right_hq,(num,num))

    # left heap에는 중간값 이하의 값만이 오게하는 과정
    if right_hq and left_hq[0][1] > right_hq[0][0]:
        left_num = heapq.heappop(left_hq)[1]
        right_num = heapq.heappop(right_hq)[0]
        heapq.heappush(left_hq,(-right_num,right_num))
        heapq.heappush(right_hq,(left_num,left_num))
        
    print(left_hq[0][1])

🔔 본 문제를 풀기 위해서는, 중간값을 힙으로 어떻게 구할 것인가에 대한 생각이 필요하다.

left_hq에는 중간값 이하의 수를, right_hq에는 중간값보다 무조건 큰 수를 위치하게 함으로써 해결할 수 있다. 왜냐하면 개수가 짝수인 경우에는, 중간에 위치한 두 숫자 중 작은 값이 중간값이 되므로 left_hq에 위치하게 된다.

1. 두 힙의 길이가 같은 경우, left_hq에 값을 추가한다.

• left_hq는 최대힙으로 구성한다. 왜냐하면 중간값 이하의 숫자들 중에서 가장 큰 값이 중간값이 되기 때문이다.

2. 길이가 다른 경우, right_hq에 값을 추가한다.

• right_hq는 최소힙으로 구성한다.

• 중간값보다 큰 숫자들 중에서 가장 작은 값이, 중간값과 가장 가까운 수이기 때문에 비교를 행할 수 있다.

3. 원소 추가를 마친 후에, left_hq에서 가장 큰 수와 right_hq에서 가장 작은 숫자를 비교한다.

• 만약 전자가 더 크다면, 두 원소의 위치를 바꾸어주는 작업을 한다.

• left_hq는 최대힙이기 때문에, -를 붙여서 추가해준다.

4. left_hq에서 가장 큰 수, 즉 중간값을 출력한다.

느낀 점 & 배운 점

1. 처음에는 하나의 힙 큐로 도전해보았는데, 원소가 추가될수록 순서가 뒤바뀌어 중간값을 구하기가 힘들었다. 결국 이러한 방식은 힙 큐의 올바른 사용법이 아니기에, 서칭을 통해 문제를 해결할 수 있었다.

2. 가장 큰 수가 처음에 위치하는 최대 힙, 가장 작은 수가 처음에 위치하는 최소 힙을 활용하여 이렇게 중간값을 구한다는 것이 신박하고 재미있는 문제였다!