가운데를 말해요

1655번: 가운데를 말해요

백준 온라인 저지 1655번 가운데를 말해요 문제입니다.

N개의 숫자를 차례차례 입력했을 때, 각각의 경우에서 중앙값을 출력하는 문제입니다.

매번 최소값이나 최대값을 출력하는 것이 아니라, 중앙값을 출력해야 하기 때문에 하나의 힙을 사용해서 풀이할 수는 없습니다. 그렇게 되면 매번 N//2의 숫자를 탐색해야 해서 시간 초과를 받게 됩니다. 

따라서, 중앙을 기준으로 하여 작거나 같은 숫자를 저장할 최대 힙과 큰 숫자들을 저장할 최소 힙을 사용하는 것이 좋습니다. 숫자가 입력되면, 힙의 크기에 따라서 두 힙 중 하나에 숫자를 넣고, 만약 두 힙이 위에 적은 기준을 만족하지 않으면 하나의 숫자를 뒤바꿔주는 방식으로 다시 힙을 교정합니다.

그 후, 중앙값을 판별하고, 출력합니다.

 

풀이 코드입니다.

from heapq import heappop, heappush
from sys import argv, stdin
input = stdin.readline


def solution(N: int, nums: list[int]) -> list[int]:
    left, right = [], []
    answer = []

    for num in nums:
        if len(left) > len(right):
            heappush(right, num)
        else:
            heappush(left, -num)

        if len(right) and -left[0] > right[0]:
            heappush(right, -heappop(left))
            heappush(left, -heappop(right))

        answer.append(-left[0])

    return answer


def main(*args: list[str]) -> int:
    N = int(input())
    nums = [int(input()) for _ in range(N)]

    answer = solution(N, nums)
    print(*answer, sep='\n')
    return 1


if __name__ == '__main__':
    main(*argv)