[백준] 1655 파이썬(python) : 가운데로 말해요 - 상세해설

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

1655번: 가운데를 말해요

 

 

import heapq
import sys # 시간초과 방지
n = int(sys.stdin.readline())
max_heap = [] #1
min_heap = []

for i in range(n):
    num = int(sys.stdin.readline())
    if len(max_heap) == len(min_heap): #2
        heapq.heappush(max_heap, -num)
    else:
        heapq.heappush(min_heap, num)
    
    if max_heap and min_heap and max_heap[0] * -1 > min_heap[0]: #3
        max_value = heapq.heappop(max_heap) * -1
        min_value = heapq.heappop(min_heap)

        heapq.heappush(max_heap, min_value * -1)
        heapq.heappush(min_heap, max_value)
    print(-max_heap[0])

 

이 문제는 시간도 촉박하고, 처음이나 끝이 아닌 중간값을 구하라 하기 때문에 다른 방법이 필요하다.

 

그 방법은 힙을 2개로 나누어서 값을 입력받는것인데 그렇게 하는 이유는 중간값을 O(1)의 속도로 받아내기 위해서다.

 

예를 들어 { 1, 2, 5, 6, 7, 11, 19 } 가 있다고 해보자. 중간값은 6이다. 하지만 파이썬에는 중간값을 O(1)의 속도로 받아주는 함수가 없다. 그래서 힙을 최대 힙과 최소 힙 2개로 나눠서 값을 입력 받는다. (절반으로 나눠서 작은 값들은 최대 힙, 큰 값들은 최소 힙에 들어가게 한다)

 

원소의 개수가 홀수 개이다. 절반으로 나눠 떨어지지 않는다. 그렇다면 2개로 나눌때 중간값은 최대와 최소 힙 중 어디에 들어가야 할까? 최대 힙에 들어가야 한다. 왜냐하면 문제에서 백준이가 외친 개수가 짝수개라면 중간에 있는 두 수 중 작은 수가 중간값이기 때문이다. 그래서 작은 원소가 들어가는 최대 힙에 중간값이 들어가야한다. 

 

#2 : 최소 힙과 최대 힙의 길이가 같으면 왼쪽인 최대 힙에 원소를 -1 곱해서 입력해준다, 길이가 다르면 균형을 맞추기 위해서 최소 힙에 입력해준다.

 

#3 : 최대 힙과 최소 힙의 길이가 0이 아니고 최대 힙의 첫번째 원소(중간값)이 최소 힙의 첫번째 원소보다 크면 그 둘의 바꿔 준다. (바꿀때 최소 힙에 -1을 곱해서 넣어준다.)

(처음 힙에 원소를 입력할때는 길이균형을 맞추기 위해 무작위로 넣었지만 최대 힙에는 중간값을 포함한 작은 원소들을 최소 힙에는 중간값이후의 큰 원소들을 넣기 위해서 if문에 걸릴때 둘을 바꿔주면 중간값보다 큰 수가 최대 힙에 들어가는것을 막아준다.)

밑에 예시가 있다

 

마지막으로 최대 힙의 첫번째 요소(중간값)을 출력한다.

입력으로 들어오는 수의 개수가 홀수개이면 최대힙에 원소의 개수가 1개 더 많게 되므로 최대 힙의 첫번쨰 요소를 출력하고, 입력으로 들어오는 수의 개수가 짝수개일때는 두 중간값 중 작은 수를 출력하라 했으므로 최대 힙의 첫번째 요소가 해당된다.

 

 

그림으로 표현하면 다음과 같다. (예제 입력 1을 예시로 들었다) 

원소가 추가될때마다 변화를 그린 그림

그림에서의 표현이고 내부에선 인덱스 0을 빼고는 최대, 최소값이 정렬되지 않는다.

최대 힙에 4개 원소가 내림차순으로(힙 내부에선 -1 곱해준 상태) 최소힙에 3개 원소가 오름차순으로 정렬되어 있다.

최대 힙의 마지막 원소에서부터 첫번째 원소까지가 절반 그리고 최소 힙의 처음부터 마지막원소까지가 절반으로 둘을 합치면 정렬된 리스트가 나오게 된다.

 

#3의 예시)

 결국 max_heap에 들어가는건 처음에 음수를 가지고 있는 숫자들이다.