[백준 1655] 가운데를 말해요

문제 출처 : https://www.acmicpc.net/problem/1655

1655번: 가운데를 말해요

우선순위 큐를 활용한 재미있는 문제입니다. 기초 알고리즘을 소개하는 다빈치 코딩 알고리즘(https://wikidocs.net/book/10280)에 등록하기에는 고난이도 문제라 판단되어 블로그에 등록합니다.

 

데이터가 입력될 때마다 배열의 중간 값을 찾는 문제 입니다. 배열을 만들어 정렬된 위치에 데이터를 넣는다면 중간값은 쉽게 찾을 수 있습니다. 하지만 배열이 커지면 커질수록 중간에 데이터를 넣는게 쉽지 않습니다. 왜냐하면 데이터가 들어올 때마다 그보다 큰 데이터는 모두 한 칸씩 뒤로 밀려나기 때문에 처리하는데 시간이 오래걸립니다.

이 문제를 해결하기 위해 우선순위 큐를 사용합니다. 우선순위 큐는 내부적으로 트리 형태이기 때문에 배열처럼 중간에 데이터가 들어와도 정렬하는데 시간이 오래 걸리지 않습니다. 문제는 우선순위 큐는 최대값이나 최소값을 우선순위를 가지고 출력합니다. 따라서 여기의 문제처럼 중간 값을 처리할 수 없습니다.

우리는 우선순위 큐 두 개를 사용하여 중간의 값을 출력하는 큐를 만들 것입니다. 한쪽의 큐에는 최댓값을 출력하는 우선순위 큐를, 다른 한쪽의 큐에는 최솟값을 출력하는 우선순위 큐를 만들어 적절하게 데이터를 추가한다면 가운데 값을 출력하는 우선순위 큐를 만들 수 있습니다.

알고리즘

가운데값을 가져오기 위한 우선순위 큐의 모습은 아래와 같습니다. 왼쪽 큐는 내림 차순 정렬하여 가장 큰 값이 노란색에 오게 됩니다. pop 함수를 실행할 때 가장 큰 숫자가 출력 됩니다. 오른쪽 큐는 오름 차순 정렬이 됩니다. pop을 하면 가장 작은 숫자가 출력됩니다. 1부터 8까지 숫자가 있으면 아래와 같은 모습이 됩니다.

그럼 예제를 통해 어떤 동작을 하게 되는지 알아보겠습니다. 먼저 빈 큐에 첫 번째 숫자 1이 들어옵니다. 큐안에 들어가는 숫자들의 균형을 맞추기 위해 양쪽 큐의 길이가 같으면 왼쪽 큐에 넣습니다.

이 때 중간값은 왼쪽 큐에 노란색에 있는 1이 됩니다. 다음으로 5를 입력 하겠습니다.

큐의 길이가 같지 않으면 오른쪽 큐에 넣습니다. 그리고 왼쪽 숫자 1과 오른쪽 숫자 5를 비교하여 순서가 맞는지 확인합니다. 오름차순으로 되어 있기 때문에 다음으로 넘어갑니다. 가운데값은 아직까지 1입니다. 다음 숫자 2를 추가하겠습니다.

양쪽 큐의 길이가 1로 같기 때문에 왼쪽 큐에 숫자를 추가합니다. 오름차순으로 정렬되어 2가 가운데 숫자가 됩니다. 2와 5를 비교하여도 순서가 맞기 때문에 바뀌지 않습니다. 다음으로 10을 추가합니다.

길이가 다르기 때문에 오른쪽 큐에 숫자를 입력합니다. 오른쪽 큐는 오름차순으로 정렬되기 때문에 10은 뒤로가게되고 가운데 숫자는 변함없이 2가 됩니다. 이번에는 예제와 다르게 99를 입력해 보겠습니다.

양쪽의 숫자가 같기 때문에 왼쪽에 99가 들어갑니다. 이제 99와 5를 비교하여 순서를 보니 잘못되어 있습니다. 99와 5를 빼서 서로 반대의 큐에 넣어줍니다.

99와 5의 위치가 바뀌어 5가 가운데 숫자가 되었습니다. 이번에는 4를 입력해 보겠습니다.

양쪽의 숫자가 다르기 때문에 4는 오른쪽에 들어가게 됩니다. 5와 4의 크기를 비교해보니 순서가 맞지 않습니다. 두 수를 꺼내어 서로 바꿔줍니다.

이렇게 4가 가운데 숫자가 되었습니다. 이렇게 우선순위 큐를 사용하여 가운데 숫자를 찾는 방법을 알아보았습니다. 해당 코드를 직접 작성해 보도록 하겠습니다.

입력받기

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

N = int(input())

left = []
right = []

입력이 많기 때문에 속도 이슈가 발생할 것을 예상하여 readline을 사용하여 입력 받는 속도를 높여주었습니다. 다음으로 우선순위 큐를 사용하기 위해 heapq의 heappop, heappush 함수를 사용할 수 있도록 import 시켜 주었습니다. 다음으로 입력 받는 숫자 N과 양쪽 큐인 left와 right를 만들어 주었습니다.

큐의 길이로 어느 큐에 넣을지 정하기

for _ in range(N):
    tmp = int(input())
    
    if len(left) == len(right):
        heappush(left, -tmp)
    else:
        heappush(right, tmp)

tmp값을 입력받아 두 큐의 길이를 비교합니다. 길이가 같으면 왼쪽 큐에, 길이가 다르면 오른쪽 큐에 넣습니다. 이 때 왼쪽 큐에 넣을 때에는 내림차순 정렬을 하기 위해 음수를 곱해주었습니다. 왼쪽 큐는 값을 꺼내올 때에도 음수를 곱해야 원래의 숫자로 돌아온다는 것을 기억해야 합니다.

큐의 값 비교하여 반대로 넣어주기

		if left and right and right[0] < -left[0]:
        left_val = heappop(left)
        right_val = heappop(right)

        heappush(left, -right_val)
        heappush(right, -left_val)
    
    print(-left[0])

먼저 두 큐가 비어있는지 확인합니다. 큐가 비어있는데 비교를 하면 에러가 발생하기 때문에 큐가 비어있는지 확인은 필수입니다. 다음으로 오른쪽 값보다 왼쪽값이 큰지 확인합니다. 왼쪽 값을 꺼낼 때에는 음수를 곱해주어야 합니다. 두 값을 비교하여 왼쪽이 더 크다면 우선순위 큐에서 꺼내어 서로 반대방향에 넣어줍니다. 이 때 오른쪽 값이 왼쪽에 들어갈 때는 내림차순을 만들기 위해 음수를 곱해주었습니다. 왼쪽 값 역시 큐에서 나왔을 때 음수이므로 음수를 곱해 양수로 만들어 오른쪽 큐에 넣어주어야 합니다. 그렇기에 왼쪽, 오른쪽 큐 모두 음수를 곱해준 것입니다. 마지막으로 왼쪽 큐의 첫 번째 값이 음수이므로 여기에 음수를 곱해 출력하면 가운데 값을 얻을 수 있습니다.

전체 코드

전체 코드를 확인해 보겠습니다.

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

N = int(input())

left = []
right = []

for _ in range(N):
    tmp = int(input())
    
    if len(left) == len(right):
        heappush(left, -tmp)
    else:
        heappush(right, tmp)
    
    if left and right and right[0] < -left[0]:
        left_val = heappop(left)
        right_val = heappop(right)

        heappush(left, -right_val)
        heappush(right, -left_val)
    
    print(-left[0])