[Python 3] BOJ - 1655 "가운데를 말해요"

 # 문제 링크 : https://www.acmicpc.net/problem/1655

1655번: 가운데를 말해요

 # 풀이 :

 이 문제의 핵심은 중앙값의 특성을 파악하는 것이다. 결론적으로는 최대 힙과 최소 힙의 이용인데 중앙값의 특성을 이용하면 왜 최대힙과 최소힙을 써야하는지 알 수 있다.

<그림 1>

 그림 1에서는 홀수일 때의 경우와 짝수일 때의 경우를 나타내었다. 홀수인 경우 중앙값은 3이다. 짝수인 경우 중앙값은 문제의 조건에 따라 2와 3중 더 작은 수라고 했기 때문에 2이다. 그런데 홀수인 경우 3을 기준으로 1,2와 4,5 양쪽을 각각의 그룹으로 나눈다고 하면 왼쪽 그룹(1,2)과 오른쪽 그룹(4,5)의 갯수는 총 2개이다. 그러나 짝수에는 중앙값의 두 후보중 하나만 중앙값이기 때문에 홀수의 경우와는 반대로 중앙값을 포함한 두 그룹의 원소 갯수가 같아야 한다. 즉, 짝수인 경우에 중앙값이 2이고 이를 왼쪽 그룹에 넣는다면 왼쪽 그룹에는 1과 2가 있으므로 2개의 원소가 있고 마찬가지로 오른쪽 그룹은 3과 4가 있으므로 2개의 원소가 있다. 

 또한 홀수의 경우 3을 만일 왼쪽 그룹에 포함한다면 3은 왼쪽 그룹 중 가장 큰 수이며 오른쪽 그룹에 포함시킨다면 오른쪽 그룹 중 가장 작은 수이다. 이 특성은 짝수일 때도 성립한다. 2는 왼쪽 그룹(1)을 기준으로 현재 왼쪽 그룹에 있는 어떤 수보다 가장 큰 수이며 반대로 오른쪽 그룹(3,4)을 기준으로 현재 오른쪽 그룹에 있는 어떤 수보다 가장 작은 수이다. 따라서 중앙값을 왼쪽 그룹에 항상 넣거나 오른쪽 그룹에 항상 넣으면 이를 기준으로 새로운 값이 들어올 때 새로운 중앙값을 처리할 수 있다.

 결국 이러한 중앙값의 특성에 따라 어떤 그룹에서 가장 작은수, 가장 큰 수만 필요하다는 것을 알 수 있고 정확히 이것이 최대 힙, 최소 힙의 의미와 동일하므로 정답을 해결하기 위해 최대 힙과 최소 힙을 모두 사용해야 한다.

<그림 2. 홀수일 때 새로운 수의 경우 중앙값 갱신>

   그런데 문제에서 입력이 반드시 정렬된 순서로 주어진다는 보장이 없고 또한 현재까지 중앙값이 다시 입력으로 주어질 수 있으므로, 즉 그림 2처럼 이미 중앙값이 3이지만 현재 들어오는 수가 다시 3일수도 있다. 그렇다면 현재 들어오는 수, 즉 now는 최소 힙과 최대 힙중 어느 곳에 넣어야 할까?

 최대 힙으로 넣는 경우를 생각하면 이는 옳지 않다. 왜냐하면 현재 now를 고려하면 홀수개에서 짝수개로 변했기 때문에 now를 최대 힙에 넣는다면 짝수의 갯수조건(왼쪽 그룹과 오른쪽 그룹의 원소의 갯수가 같아야함)을 만족시키지 못하기 때문이다. 따라서 now는 최소 힙에 넣어야 함을 알 수 있다.

<그림 3. now를 최소힙에 넣는 경우>

 따라서 그림 3과 같이 now를 최소 힙에 넣는다. 그런데 짝수의 대소조건(두 중앙값 후보 중 작은 수를 중앙값으로 함)을 비교해야 하나 최소 힙의 top과 최대 힙의 top이 같으므로 둘 중 어느 수를 중앙값으로 해도 해당 경우는 정답을 구할 수 있다.

 그러나 만일 now가 3이 아니라 2였다면 최소 힙의 top이 중앙값이 되어야 한다. 그런데 원래 1,2,3,4,5의 수열이 정렬한 후 현재 1,2,2,3,4,5가 된다. 이런 경우에서 2는 반드시 왼쪽 그룹, 즉 최대 힙에 있어야 중앙값의 조건을 만족하는데 현재 최소 힙에 있기 때문에 조건을 만족하지 않는다. 따라서 갯수조건을 먼저 파악한 후 각 힙의 top을 비교하여 최소 힙의 top이 최대 합의 top보다 작다면 서로를 바꿔주면 올바르게 중앙값을 구할 수 있다.

<그림 4. 짝수일 때 새로운 수의 경우 중앙값 갱신>

 그러나 지금까지는 홀수인 경우에 새로운 수가 들어올 때 중앙값의 적절한 갱신만을 생각했다. 따라서 짝수인 경우에 새로운 수가 들어올 때 중앙값의 적절한 갱신을 고려한다면 문제의 정답을 구할 수 있다. 홀수의 갯수 조건은 중앙값을 기준으로 왼쪽, 오른쪽 그룹으로 나누면 두 그룹의 원소의 갯수는 같다는 것인데 중앙값을 어떤 그룹이던지 한쪽 그룹에만 추가한다면 항상 두 그룹의 원소의 갯수는 1만 차이가 나야 한다. 따라서 now를 최소 힙에 넣든 최대 힙에 넣든 현재는 짝수지만 now가 들어옴으로 인해서 전체 갯수가 홀수가 되었기 때문에 어디로 넣든 홀수일 때의 갯수 조건은 만족한다.

<그림 5. now를 최대힙에 넣는 경우>

 now를 포함한 후 정렬하게 되면 1,2,3,4,7이 되므로 7은 최소 힙에 들어가야 함을 알 수 있다. 왜냐하면 최대 힙에 있게 되면 1,2,7,3,4가 되어야 하기 때문이다. 그러나 최소 힙에 들어가면 7은 최소 힙에 있는 어떤 수보다도 가장 크기 때문에 최소 힙의 가장 아래에 위치하게 된다. 따라서 1,2,3,4,7의 구조를 따르게 된다.

 그런데 짝수의 대소 조건을 생각하게 되면 최소 힙의 top이 최대 힙의 top보다 작다면 두 수를 바꿔주는 것이었다. 이것을 현재 그림 5의 경우에 적용하게 되면 7은 최소 힙의 어떤 수보다도 가장 큰 수이므로 7은 최소 힙의 최하단에 위치하게 되고 홀수일 때 대소조건을 만족하기 때문에 짝수의 대소조건을 홀수의 경우에도 적용해도 홀수의 경우 중앙값을 구하는데 문제가 없다. 따라서 짝수인지 홀수인지에 상관없이 정답을 구할 수 있다.

 

 # 코드

import sys, heapq

# max_heap : 최대 힙
# min_heap : 최소 힙
max_heap = []
min_heap = []

# ans : 정답 저장 배열
ans = []

# 수열의 갯수 입력
n = int(sys.stdin.readline())
for i in range(n):

    # 수열의 원소 입력
    m = int(sys.stdin.readline())
    
    # 같은 개수라면 무조건 최대힙으로 넣는다
    if len(max_heap) == len(min_heap):
        heapq.heappush(max_heap, (-1 * m,m))
        
    # 개수가 다르다는 것은 이전에 최대힙에 넣은 것이기 때문에
    # 개수를 맞춰주기 위해 최소힙에 넣는다
    else:
        heapq.heappush(min_heap, (m,m))
        
    # 앞에서 갯수 조건을 맞춰주었기에 대소 조건을 비교한다
    # 최소 힙의 top이 최대 힙의 top보다 작다면 두수를 바꿔준다
    if min_heap and max_heap[0][1] > min_heap[0][1]:
        temp_min = heapq.heappop(min_heap)[1]
        temp_max = heapq.heappop(max_heap)[1]
        heapq.heappush(max_heap, (-1 * temp_min, temp_min))
        heapq.heappush(min_heap, (temp_max, temp_max))
        
    # 갯수조건, 대소조건을 만족하였으니 최대 힙의 top이 중앙값이다
    ans.append(max_heap[0][1])

# 정답 출력
for i in ans:
    print(i)