[백준 1655 파이썬] 가운데를 말해요 (골드2, 우선순위 큐)

알고리즘 유형 : 우선순위 큐
풀이 참고 없이 스스로 풀었나요? : X

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

소스 코드(파이썬)

import sys
import heapq
input = sys.stdin.readline

left_heap = []
right_heap = []

for _ in range(int(input())):
    n = int(input())
    
    if len(left_heap) == len(right_heap):
        heapq.heappush(left_heap, (-n, n))
    else:
        heapq.heappush(right_heap, n)
        
    if right_heap and left_heap[0][1] > right_heap[0]:
        smaller = heapq.heappop(right_heap)
        bigger = heapq.heappop(left_heap)[1]
        
        heapq.heappush(left_heap, (-smaller, smaller))
        heapq.heappush(right_heap, bigger)
        
    print(left_heap[0][1])

풀이 요약

1. 이 문제의 핵심 아이디어는 최대 힙, 최소 힙 두 개를 사용하는 것이다.

2. 최대 힙, 최소 힙 두개를 사용하는 이유

   테스트케이스의 숫자를 계속 리스트에 넣고, 오름차순으로 생각하면서 중간값이 몇 번째에 있는지 계산해보자.

   
   

   [1] : 길이 1, 중간값 1째
   [1, 5] : 길이 2, 중간값 1번째
   [1, 2, 5] : 길이 3, 중간값 2번째
   [1, 2, 5, 10] : 길이 4, 중간값 2번째
   [-99, 1, 2, 5, 10] : 길이 5, 중간값 3번째
   [-99, 1, 2, 5, 7, 10] : 길이 6, 중간값 3번째
   [-99, 1, 2, 5, 5, 7, 10] : 길이 7, 중간값 4번째

   
   

   잘 보면 규칙이 보인다. 리스트의 길이에 대해, 중간값은 (길이+1)//2 번째에 있다.

   이제 우리는 최대 힙 left_heap과 최소 힙 right_heap을 만들 것이다. left_heap의 모든 원소는 right_heap의 원소보다 작거나 같다. 이름 그대로 좌우(대소) 관계를 만족하는 것이다.

   그리고 길이와 중간값 위치 규칙을 활용하여, left_heap의 루트 노드, 즉 left_heap에서 가장 큰 수를 뽑았을 때 그 값이 중간값이 되게끔 입력값들을 추가해보자.

   
   1) len(left_heap) == len(right_heap)인 경우

   여기서 입력값 하나를 힙에 넣으면 총 개수는 홀수개가 된다. 이 개수를 k라고 해보자. 그럼 중간값은 (k+1)//2번째 수에 있고, 개수가 홀수니까 더 단순히 생각하면 k//2 + 1 번째에 있다.(예를 들어 7개면 (7//2) + 1 = 4번 째에 중간값이 있음)

   그런데 이 때 오름차순 기준으로 처음부터 k//2 + 1 개의 수가 left_heap에 있고, k//2 개의 수가 right_heap에 있다. (직접 대입해보면 이해가 될 것이다.)

   그러므로, 중간값이 left_heap에 있게 되고, 이 때 중간값의 위치 서수와 left_heap의 길이가 같으니, left_heap에서 가장 큰 수, 즉 left_heap의 루트 노드가 곧 중간값이게 된다.

   
   

   2) len(left_heap) != len(right_heap)인 경우

   이 때는 위의 경우를 거쳐 left_heap에 수의 개수가 right_heap보다 하나 더 많은 상태이다. 이 때 right_heap에 입력값을 넣어주면, 총 개수는 짝수개(k)가 되고, 이 때의 중간값은 모든 수를 오름차순 상태로 두었을 때 기준으로 k//2 번째에 있다. 그런데 이 때 left_heap에 들어 있는 수가 k//2개이므로, left_heap에서 가장 큰수, 그러니까 루트 노드가 중간값이다.

   이로써 입력 값을 넣을 때마다 항상 left_heap의 루트 노드가 중간값이게 되었으므로, 각 단계에서 left_heap의 루트 노드를 인덱싱(리스트이므로)하여 출력해주면 된다.

3. 하지만 주의해야할 점이 있다. 위에서 설명한 로직에 따라 입력값을 어떤 힙에 넣긴 넣었는데, 이 후 최대 힙의 모든 원소가 최소 힙의 모든 원소보다 작거나 같아야한다는 조건을 만족하지 않게 되어버리는 경우의 수가 있는데, 이건 어떻게 해결할까?

   예를 들어, left_heap = [5, 3, 1] 이고, right_heap = [7, 9, 11] 이라고 가정해보자.

   이 때 다음 입력값 10을 넣을 때는, 위에서 말한 "총 개수"에 따른 조건문을 따라, 두 힙의 길이가 같으므로 left_heap에 넣게 된다.

   그런데 앞서 left_heap의 모든 원소는 right_heap의 모든 원소보다 작아야한다고 했다. 10을 left_heap에 넣어버리면, 이 것은 right_heap의 7, 9보다 크므로 이 규칙에 위배된다.

   그래서, 입력값을 넣어줄 때마다, 일단 넣고 난 다음에, 두 힙의 루트 노드의 대소 관계를 비교하여, 조건에 위배된다면 그 수를 서로 바꿔줘버리면 된다.

   
   

   예를 들면, 총 개수 조건 로직에 따라 10을 left_heap에 일단 넣어보자. 그러면

   left_heap = [10, 5, 3, 1]
   right_heap = [7, 9, 11]

   이렇게 된다.

   이 때 left_heap의 루트 노드인 10과 right_heap의 루트 노드인 7을 비교하면, left_heap의 루트 노드가 더 크므로, 일단 두 루트 노드를 각각 빼내준 뒤, 10을 right_heap에 넣어주고 7을 left_heap에 넣어준다. 그러면

   left_heap = [7, 5, 3, 1]
   right_heap = [9, 10, 11]

   이제 대소 관계도 만족하고 left_heap의 루트 노드가 중간값인 것도 만족한다.

   왜 두 힙의 루트 노드만 비교하면 되는지 이해가 안 될 수도 있는데, 입력값을 넣기 전의 두 힙은 이미 이 전 과정을 거치면서 대소 관계, 중간값 위치 조건을 모두 만족하는 상태이고, 더불어 입력값은 단 하나씩만 넣어 주기 때문에 두 힙의 "루트 노드"만 비교해주면 되는 것이다.

   좀 더 쉽게 설명하면, 본디 right_heap에 들어가야 로직에 맞게 되는 입력값을 left_heap에 넣었을 때, right_heap은 left_heap에 들어간 원소(left_heap의 루트 노드에 위치해 있음)를 그대로 가져오면서, 대신 자신의 가장 작은 값인 루트 노드를 left_heap에 줘 버리면, 총 개수 로직도 만족을 하고, 대소 관계 로직도 만족을 하게 된다.

   그리고 이 것은 left_heap에 들어가야 맞는 입력값을 right_heap에 넣어버리는 경우에도 같은 원리가 적용된다.

   테스트 케이스를 가지고 직접 공책에 적어보면 이해가 잘 될 것이다.

배운 점, 어려웠던 점

• 수의 총 개수에 따라, 오름차순 기준으로 중간값이 어디에 위치해있는지에 대한 규칙은 찾았는데, 그걸 최대 힙, 최소 힙 두개를 만들어서, 중간값이 항상 최대 힙의 루트 노드에 있게끔 입력값을 매번 넣어주는 아이디어는 정말 상상도 못했다...

• 입력 값을 최대 힙 또는 최소 힙에 넣을 때, 넣고 나서 최대 힙의 모든 원소가 최소 힙의 원소보다 작거나 같아야 한다는 조건을 만족시키기 위해, 모든 원소를 비교하지 않고 두 힙의 "루트 노드"만 비교해도 되는 이유를 이해하는 것에서 조금 애먹었다.