백준 / 골드 2 / 1655. 가운데를 말해요

맨날 짤 선정 능력만 늘고 알고리즘은 안 늘고... 반성하는 중입니다.

생각해봅시다!

• 지금까지 입력된 수들 중 중간값을 매번 출력하면 됩니다.
• 다만 시간제한이 0.6초고, 수의 최대 개수가 $100,000$개인 만큼, 정렬로 풀면 시간초과 납니다.
• 조금 더 빠르고 신박한 방법이 필요한데... 저장한 숫자들 중에서 중간값을 바로 꺼낼 수 있는 그런 자료구조가 없을까요?
• 생각해 보니, 최솟값이나 최댓값을 바로 꺼낼 수 있는 자료구조들은 있었죠. 최대 / 최소 힙이였죠? 뭔가 이거 두 개를 응용해 볼 순 있겠네요.
• 사실 직관적으로 떠올리기 쉽지 않아요 ㅠㅠ 흐엉 나 울어

최대 힙, 최소 힙 사용하기

• 우리는 중간값을 바로 반환받고 싶습니다.
• 최대 힙이, 중간값을 반환하게 하면 어떨까요?
  • 중간값 및 중간값보다 작은 숫자들을 최대 힙에 저장하면, 숫자들 중 최댓값인 중간값을 항상 확인 및 반환할 수 있습니다.
• 중간값보다 큰 숫자들은 최소 힙에 저장합니다.
  • 가끔 최대 힙에 중간값보다 큰 숫자가 저장되어 있거나, 중간값이 저장되어 있지 않는 경우가 있습니다.
  • 이럴 때, 중간값보다 큰 숫자를 받아오거나, 중간값을 가져올 때 최소 힙을 사용합니다.
  • 뭔 소린지 모르겠다고요? 저도 그랬어요 무하하 일단 읽어보십쇼.

값 푸시하는 방법

• 현재 중간값 이하인 숫자들은 최대힙에, 큰 숫자들은 최소힙에 저장합니다.
  • 첫 숫자의 경우 아직 중간값이 없는데, 중간값으로 만들어주기 위해 최대힙에 저장합니다.

• 그런데 현재까지 입력된 숫자가 $A$개일 때, 최대힙에는 총 (A + 1) // 2개의 원소가 있어야 합니다.
  • 첫 숫자부터 중간값까지의 숫자 수가 (A + 1) // 2개가 됩니다.
  • e.g., 3개인 경우, 2번째 숫자가 중간값이므로 (3 + 1) // 2 = 2
  • e.g., 6개인 경우, 6번째 숫자가 중간값이므로 (6 + 1) // 2 = 3
• 최대힙의 원소가 (A + 1) // 2개보다 적은 경우
  • 최소힙에서 원소를 팝해 최대힙으로 가져옵니다
• 최대힙의 원소가 (A + 1) // 2개보다 많은 경우
  • 최대힙에서 원소를 팝해 최소힙으로 푸시합니다

• 아무튼 푸시를 한 뒤, (필요할 시) 힙 간에 원소를 옮긴 후, 중간값을 최대힙의 최댓값으로 갱신합니다.

• 이런 과정을 거치면 이렇게 매번 중간값이 갱신됩니다.

풀이

import sys
import heapq
import math
input = sys.stdin.readline

min_heap = []	# 중앙값 이후 값을 담음
max_heap = []	# 중앙값 및 중앙값 이전 값을 담음

N = int(input())  # 총 숫자 수

for i in range(N):
    num = int(input())
    if i == 0: # 첫번째 숫자는, 최대 힙에 넣기
        heapq.heappush(max_heap, -num)
        mid_value = num

    else:
        if num > mid_value: # 중앙값 이후의 값
            # 최소 힙에 푸시
            heapq.heappush(min_heap, num)
        
        else: # 중앙값 이전의 값
            # 최대 힙에 푸시: 푸시할 때 음수 처리
            heapq.heappush(max_heap, -num)
    
        # 현재 최대 힙에 있어야 하는 원소 수
        max_heap_len = (i + 2) // 2
        
        # 최소 힙에서 최대 힙으로 원소 보내기
        if len(max_heap) < max_heap_len:
            popped = heapq.heappop(min_heap)
            heapq.heappush(max_heap, -popped) # 최대힙은 음수처리!
        # 최대 힙에서 최소 힙으로 원소 보내기
        elif len(max_heap) > max_heap_len:
            popped = -heapq.heappop(max_heap) # 최대힙은 음수처리!
            heapq.heappush(min_heap, popped)
    
    # 중앙값은 최대 힙의 루트 노드
    mid_value = -max_heap[0]
    print(mid_value)

• 최대 힙에 푸시할 땐 heapq.heappush(max_heap, -num) 식으로 푸시 값을 음수 처리하고, 팝할 땐 반환값을 -heapq.heappop(max_heap), -max_heap[0]으로 음수 처리해야한다는 점 기억합시다.
  • 파이썬의 heapq 모듈은 최소 힙 기반이므로, 이런 처리가 필요합니다.
• 여기선 인덱스가 i일 때 i + 1개 숫자가 입력된 상태므로, 최대 힙에 ((i + 1) + 1) // 2 -> (i + 2) // 2개의 원소가 있어야 정상입니다.

시간 복잡도

• 1부터 $N$까지 수에 대해...
• heappush heappop 모두 $O(\log N)$
• $O(N \log N)$. 이정도면 시간제한 통과해요~

기억할 점

최대 힙, 최소 힙을 둘 다 활용하면 최솟값, 최댓값 아닌 수도 꺼내올 수 있다.

단지 그 방법을 떠올리는 건 너의 몫이라는 점 ^^