[힙] 백준 #2696 중앙값 구하기 - Python

🗂️ 문제

링크: https://www.acmicpc.net/problem/2696

2696번: 중앙값 구하기

 

💡 접근법

가운데를 말해요 문제와 비슷한 문제이다. 즉, 최대 힙과 최소 힙을 생성하여 중앙값을 구하면 된다.

leftHeap에는 중간값보다 같거나 작은 값들이 저장되고, rightHeap에는 중간값보다 큰 값들이 저장된다.

 

leftHeap의 루트를 중간값으로 꺼낼 것이다. 이때 만약 수의 개수가 짝수 개라면 중간에 있는 두 수 중 작은 수를 말해야 한다는 문제 조건이 있다. 따라서 rightHeap과 leftHeap의 길이가 같을 경우에는 leftHeap에 새로운 값을 push해야 한다.

 

또한 만약 leftHeap의 루트가 rightHeap의 루트보다 클 경우에는 두 값을 swap해줘야 한다. leftHeap의 루트는 중간값이다. 또한 leftHeap에는 중간값보다 작은 값들이 저장되어야 한다. 그런데 leftHeap의 루트가 rightHeap의 루트보다 크다면 중간값보다 작은 값이 rightHeap에 저장되어 있는 것이기 때문이다.

 

이 외의 문제 조건에 맞게 구현해야 한다.

• 홀수 번째 수에 대해서만 중앙값을 구해야 하므로 answer이라는 리스트를 선언하여 구한 중앙값들에 대해 저장한다. 이때 리스트의 인덱스는 0부터 시작하기 때문에 (i+1) % 2 == 1의 조건을 만족하는 인덱스에 대해 중앙값을 구한다.
• 입력으로 한 줄에 원소가 10개씩 주어지고, 출력 시에도 한 줄에 10개씩 출력해야 한다. 따라서 나머지 연산을 이용한다.

 

😎 내 코드

import sys
import heapq

T = int(sys.stdin.readline())
for _ in range(T):
    M = int(sys.stdin.readline())
    
    left_heap = [] #최대 힙
    right_heap = [] # 최소 힙
    answer = []

    arr = list(map(int, sys.stdin.readline().split()))
    for i in range(M//10):
        arr += list(map(int, sys.stdin.readline().split()))

    for i, a in enumerate(arr):
        if len(left_heap) == len(right_heap):
            heapq.heappush(left_heap, -a)
        else:
            heapq.heappush(right_heap, a)
        
        if left_heap and right_heap and -left_heap[0] > right_heap[0]:
            large = heapq.heappop(left_heap)
            small = heapq.heappop(right_heap)
            heapq.heappush(left_heap, -small)
            heapq.heappush(right_heap, -large)
        
        if (i+1) % 2 == 1:
            answer.append(-left_heap[0])
    
    print(len(answer))
    for i, ans in enumerate(answer):
        if i % 10 == 0 and i > 1:
            print()
        print(ans, end=' ')
        
        if i == len(answer)-1:
            print()

 

🧐 배운 점

• 리스트 이어붙이기
  리스트 a = ['BlockDMask', 333]와 리스트 b = [1,2,3]을 더하면 하나의 또 다른 리스트인 ['BlockDMask', 333, 1, 2, 3]이 생성된다.

 

✅ 정답 확인