[백준] 1655. 가운데를 말해요/Python - Gold2

❓문제

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

성능 요약

메모리: 39016 KB, 시간: 228 ms

분류

자료 구조, 우선순위 큐

 

문제 설명

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

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

✍🏻풀이

정말 간단하게 보면 반복문을 이용해서 값을 넣을 때마다 정렬해서 풀 수 있다. 그러나 시간복잡도가 O(n²logn)으로 시간 초과가 발생한다. 그러므로 우리는 우선순위 큐를 이용해서 풀면 된다. 다만, 이 때 최소 힙만 이동한다면 원하는 출력 값이 나오지 않는다. 따라서 최소 힙과 최대 힙을 두 개 다 이용하겠다. 왼쪽으로는 최대 힙을 사용해서 작은 수들만 모아놓고, 오른쪽에는 최소 힙을 사용해서 큰 수들만 모아놓겠다.

먼저, 짝수일 때와 홀수일 때를 고려하여 길이가 같다면 왼쪽에 넣어주고 다르다면 오른쪽에 넣어주겠다. 그런데 만약 1, 2, 5 순으로 입력을 받는다면 left에 1과 5가 들어가고 right에는 2가 들어갈 것이다. 우리는 2와 5를 빼서 바꿔주면 된다. 그러면 left는 실제로는 [-2, -1], right는 [5]가 될 것이다.

문제에서 주어진 입력값대로 넣으면 다음과 같이 된다.

그러면 우리는 left에 있는 값에 -를 붙여서 출력하면 중앙값을 도출할 수 있다.

💻코드

import sys
from heapq import heappush, heappop

input = sys.stdin.readline

N = int(input()) # 정수의 개수
left_nums = []
right_nums = []

for i in range(1, N+1):
    n = int(input())
	
    # 총 길이가 짝수일 때와 홀수일 때를 고려.
    if len(left_nums) == len(right_nums):
        heappush(left_nums, -n)
    else:
        heappush(right_nums, n)
	
    # 루트 값을 비교하여 작은 수대로 큰 수대로 모음. 정렬한다고 생각하면 됨.
    if right_nums and -left_nums[0] > right_nums[0]:
        right = heappop(right_nums)
        left = heappop(left_nums)

        heappush(right_nums, -left)
        heappush(left_nums, -right)
    print(-left_nums[0])

📍후기

나는 처음에 문제를 접근했을 때 최소 힙으로 풀면 알아서 작은 수대로 정렬 되는 것이라고 생각했다. 그러나 힙으로 넣은 후 최소힙은 최소힙대로 최솟값을 pop하고 최대힙이면 최댓값을 pop한다는 것을 알게 되었다. 이 문제를 풀며 힙에 대해서 다시 정리할 수 있었다.