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

이 문제를 차례대로 숫자가 들어오고 , 그 동안 들어온 값들 중 중위값을 구하는 문제다.

처음에는 이진정렬을 이용해 매번 정렬을 하고, 그 후 중위값을 구해야겠다고 생각했지만 이는 시간을 초과했다.

잘 모르겠어서 찾아보니 heap( 우선순위 queue)를 이용해서 푸는 문제라고 한다.

 

---

먼저 우선순위 큐는 이진트리를 이용해 정렬된 값을 가지는 자료구조를 말한다.

🌈 자료구조:: 우선순위 큐(Priority Queue) -> 정리를 잘 해놓은 사이트라 참고하면 쉽게 이해할 수 있을 것이다.

따라서 이 자료구조를 이용하면 $O(log_2 n)$으로 정렬된 자료구조를 사용할 수 있다.

우선순위 큐 예시 그림

---

이 문제의 풀이방식의 컨셉은 2개의 heap를 이용하고 왼쪽 heap의 루트노드에 중위값이 오도록 유지하는 것이다.

 

먼저 left heap은 max heap으로 root node에 가장 큰 값이 오게 되고,

right heap은 min heap으로 root node에 가장 작은 값이 오게 된다.

 

left heap에 항상 median값이 오게 하기 위해서는, left heap의 길이가 right heap의 길이보다 같거나 1이 더 길어야한다.

왜냐하면 짝수인경우 ( 두 개의 길이가 같은 경우 ) , 더 작은 값을 출력하라고 문제에서 제시했고

홀수인 경우에는 위의 예시처럼 길이가 7이면 3번째 값이 닶이 되기 때문이다.

 

이 2개의 컨셉을 이해하면 이 문제 풀이를 전부 이해한 것이다.

이를 파이썬으로 구현하면 다음과 같다.

 

import heapq
from sys import stdin, setrecursionlimit
setrecursionlimit(10**6) # 파이썬 최대 재귀깊이 높이기

stdin = open('1655/sample_input.txt', 'rt') # 테스트할 때, 제출할 때는 주석처리 해주기

n = int(stdin.readline().strip())

# python에서 제공하는 heap은 min heap. root node에 가장 작은 값이 옴.
# 따라서 max heap을 사용하기 위해 입력값을 음수로 바꿔 입력
left_heap = [] # max heap
right_heap = [] # min heap

for _ in range(n):
  input_number = int(stdin.readline().strip())
  
  # left heap의 개수가 right heap의 개수보다 항상 많거나 같아야 left heap의 root node가 median을 가짐 
  # left heap과 right heap의 개수 맞춰주기
  if (len(left_heap) == len(right_heap)):
    heapq.heappush(left_heap, -input_number)
  else:
    heapq.heappush(right_heap, input_number)
  
  # left_heap(max heap)이 right_heap(min heap)보다 클 때 -> 두 개의 위치를 바꿔준다. 
  if right_heap and -left_heap[0] > right_heap[0]:
    left_number = heapq.heappop(left_heap)
    right_number = heapq.heappop(right_heap)
    
    heapq.heappush(left_heap, -right_number)
    heapq.heappush(right_heap, -left_number)
    
  # left heap (max heap)에서 정답 꺼내기
  print(-left_heap[0])

 이 코드에서 추가된 개념은, 

파이썬에서는 min heap만 제공하기에 left heap(max heap)을 구현하기 위해 입력값이 -1을 곱해줬다는 것과 , 

left heap의 root node가 right heap의 root node보다 크면 두 개의 값을 바꿔줬다는 것이다.

( 홀수인 경우 median은 작은 값이 되어야하니까 )

 

---

heap(priority queue)를 연습하고 익힐 수 있는 좋은 문제라고 생각한다.