[Python] 백준 2696 - 중앙값 구하기 문제 풀이

Overview

BOJ 2696번 중앙값 구하기 문제 Python 풀이
분류: Data Structure (자료구조)

---

문제 페이지

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

---

풀이 코드 - heapq

from sys import stdin
from heapq import heappush, heappop
input = lambda: stdin.readline().rstrip()


def get_median() -> None:
    left, right, ans = [], [], [nums[0]]
    mid = nums[0]

    for i, num in enumerate(nums[1:], 1):
        if num < mid:
            heappush(left, -num)
        else:
            heappush(right, num)

        if i % 2 == 0:
            if len(left) > len(right):
                heappush(right, mid)
                mid = -heappop(left)
            elif len(left) < len(right):
                heappush(left, -mid)
                mid = heappop(right)
            ans.append(mid)

    print(M // 2 + 1)
    for i, num in enumerate(ans):
        if i != 0 and i % 10 == 0:
            print()
        print(num, end=' ')
    print()


if __name__ == "__main__":
    T = int(input())
    for _ in range(T):
        M = int(input())
        nums = []
        for _ in range(M // 10 + 1):
            nums += list(map(int, input().split()))

        get_median()

Heapq 두 개를 사용하여 풀이한다. left는 중앙값보다 작은 값을 저장하는 최대힙, right는 중앙값보다 큰 값을 저장하는 최소힙이다.
홀수 번째 숫자는 힙에 넣은 뒤에 중앙 값을 구한다. left와 right 중에서 길이가 긴 힙에서 중앙값을 뽑는다. 만약 길이기 같다면 중앙값은 그대로가 된다.

---

풀이 코드 2 - sort 이용

from sys import stdin
input = lambda: stdin.readline().rstrip()


if __name__ == "__main__":
   T = int(input())
   for _ in range(T):
       M = int(input())
       nums = []
       for _ in range(M // 10 + 1):
           nums += list(map(int, input().split()))

       print(M // 2 + 1)
       for i in range(1, len(nums) + 1, 2):
           print(sorted(nums[:i])[i // 2], end=' ')
       print()

heapq를 사용하지 않고, 매 번 sort를 이용하여 수열을 정렬해가며 중앙값을 구하면 어떨까?
채점 결과 통과는 했지만, 이전 풀이에 비해 시간이 오래 걸린다.