[python]2696_중앙값 구하기

1 사용한 자료구조 및 개념 : heap

 

💡 문제풀이 아이디어 및 어려웠던 점

💫 아이디어

1️⃣ heap을 사용하여 중앙값을 구한다.

2️⃣ max_heap에는 중앙값을 포함한 왼쪽 부분을 저장하고, min_heap에는 오른쪽 부분을 저장한다.

3️⃣ 배열은 0번부터 시작이니 짝수번째에서 중앙값을 구한다.

4️⃣ (배열의)짝수번째일 때, 숫자가 min_heap의 최솟값보다 크다면 min_heap에 추가하고, 작다면 max_heap에 추가한다. min_heap에 추가했다면 max_heap에서 숫자 하나를 pop하여 min_heap에 추가한다.

→ 이때, 중앙값은 정렬한 수열에서 항상 가운데 있다. 따라서 중앙값을 효율적으로 찾기 위해 max_heap의 크기는 min_heap의 크기보다 1만큼 크거나 같아야 한다.

✅ 1, 3, 5, 7, 9라는 수를 차례대로 받는다고 생각해보자.
숫자 1이 오면, max_heap에 1이 추가되며, 중앙값은 1이다.
숫자 3이 오면, max_heap에 1, min_heap에 3이 추가되며, 중앙값은 여전히 1이다.
숫자 5가 오면, max_heap에 1과 3, min_heap에 5가 추가되며, 중앙값은 3이다.
숫자 7이 오면, max_heap에 1과 3, min_heap에 5와 7이 추가되며, 중앙값은 여전히 3이다.
숫자 9가 오면, max_heap에 1, 3, 5, min_heap에 7과 9가 추가되며, 중앙값은 5이다.

5️⃣ (배열의)홀수번째 일 때, 숫자가 max_heap의 최댓값보다 작다면 max_heap에 추가하고, 크다면 min_heap에 추가한다. max_heap에 추가했다면 min_heap에서 숫자 하나를 pop하여 max_heap에 추가한다.

---

Solution Code & 주석

# heap을 사용한 코드
import sys
from heapq import heappush, heappop

input = sys.stdin.readline

for _ in range(int(input())):
    m = int(input())
    nums = []
    for _ in range(m // 10 + (m % 10 > 0)):
        nums.extend(map(int, input().split()))
    max_heap = []
    min_heap = []
    answer = []
    for i in range(m):
        if i % 2 == 0:
            if not min_heap or min_heap[0] >= nums[i]:
                heappush(max_heap, (-nums[i]))
            else:
                heappush(min_heap, nums[i])
                heappush(max_heap, -heappop(min_heap))
            answer.append(-max_heap[0])
        else:
            if -max_heap[0] <= nums[i]:
                heappush(min_heap, nums[i])
            else:
                heappush(max_heap, -nums[i])
                heappush(min_heap, -heappop(max_heap))
    out_len = len(answer)
    print(out_len)
    for i in range(out_len // 10):
        print(*answer[i * 10 : (i + 1) * 10])
    if out_len % 10 > 0:
        print(*answer[out_len // 10 * 10 :])

코드 동작 그림