백준 1655 가운데를 말해요

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

1655번: 가운데를 말해요

처음에 풀었을 때 문제에 대해서 깊게 고민하지 않았다. 하나의 배열에서 계속해서 배열의 크기에 따라 중간 index를 계산하면 되지 않나? 이렇게 생각하면 배열의 길이가 홀수 와 짝수 일때 구분하고 알고리즘을 짜면 된다. 해서 다음과 같이 코드를 작성했다.

stack = []
N = int(input())
for _ in range(N):
  stack.append(int(input()))
  stack.sort()
  if len(stack) % 2 == 0:
    mid = len(stack) // 2 -1
  else:
    mid = len(stack) // 2
  print(stack[mid])

예상한 사람들은 있겠지만 시간초과가 뜬다. 문제는 반복문에서 숫자를 입력할 때마다 stack 배열을 정렬해야 한다는 것, 시간 초과는 이때문인것으로 보인다. 이진 탐색을 사용해야 하나 싶지만 문제의 조건을 보면 숫자를 입력할 때마다 가운데 수를 말해야 한다. 그러면 기존의 Stack이나 Que같은 LIFO(Last in first out), FIFO(First in first out)구조로는 안될거 같은데? 그렇다, 그래서 특정 기준으로 원소를 뽑는 우선순위 큐인 힙을 사용한다. 힙은 원래 완전 이진 트리의 구조이지만 우리는 문제를 풀기 위해 트리 구조를 구현하는 비효율적인 방법을 추구하지 않으므로 min,max heap 배열 두개로 입력된 값을 나누어 받도록 하겠다. 설명은 문제에서 주어진 입력 예시: (7, 1, 5, 2, 10, -99, 7, 5)로 하겠다.

힙의 정확한 개념을 여기서 다루진 않겠다. min,max heap의 변화가 정확하게 저렇게 되는 이유는 개념을 정확히 알아야 이해가 가능하기에 넘어가고 일단 max_heap은 가장 큰 수를 root로 두고 min_heap은 가장 작은 수를 root로 이진트리를 전개한다는 것을 알면 된다. 그리고 그 root는 배열에서는 0번 index이다. 이를 이용하여 max_heap[0]보다 큰 수를 min_heap에 보내면 큰 수들 중 가장 작은 수가 root에 저장되고 그렇지 않은 작은 수들을 max_heap에 넣으면 max_heap[0]은 작은 수 중 가장 큰 수를 저장한다. 즉 max_heap[0]과 min_heap[0]으로 중간 값을 추적할 수 있다. 이를 코드로 구현하면 다음과 같이 된다

 

import heapq
import sys
N =  int(sys.stdin.readline())
min_heap = []
max_heap = []

for _ in range(N):
  num = int(sys.stdin.readline())

  #heap 배열 저장 논리
  #max_heap가 비어있거나 숫자가 max_heap의 가장 작은수 보다 크거나 작다면
  if not max_heap or num <= -max_heap[0]:
    heapq.heappush(max_heap, -num)
  else:
    heapq.heappush(min_heap, num)

  #min과 max heap의 크기 맞추는 코드
  if len(max_heap) > len(min_heap) + 1:
    heapq.heappush(min_heap, -heapq.heappop(max_heap))
  elif len(min_heap) > len(max_heap):
    heapq.heappush(max_heap, -heapq.heappop(min_heap))
  
  #가운데 수 mid 출력하기
  if len(max_heap) == len(min_heap):
    mid = min(-max_heap[0], min_heap[0])
  else:
    mid = -max_heap[0]
  print(mid)