중앙값

• 중앙값 측정과 동일한 문제이다. 입출력 형식, 인풋값의 범위까지 모두 동일하니 풀이는 그쪽을 참고.
• 다만, 테스트셋이 조금 더 타이트한것 같다. 처음에 실수로 버그가 있는 코드를 제출했었는데, 중앙값 측정에서는 통과되었으나 이 문제에서는 오답이 뜬 적이 있었다.

"""Solution code for "BOJ 1572. 중앙값".

- Problem link: https://www.acmicpc.net/problem/1572
- Solution link: http://www.teferi.net/ps/problems/boj/1572
"""

import sys
from teflib import priorityqueue

INF = float('inf')


def main():
    N, K = [int(x) for x in sys.stdin.readline().split()]
    nums = [int(sys.stdin.readline()) for _ in range(N)]

    median = INF
    median_sum = 0
    min_heap = priorityqueue.UpdatableHeap()
    max_heap = priorityqueue.UpdatableHeap()
    for i, num in enumerate(nums):
        if i >= K:
            num_to_delete = nums[i - K]
            if num_to_delete <= median:
                max_heap.delete(-num_to_delete)
            else:
                min_heap.delete(num_to_delete)
        if num <= median:
            max_heap.push(-num)
            while len(max_heap) > len(min_heap) + 1:
                min_heap.push(-max_heap.pop())
        else:
            min_heap.push(num)
            while len(max_heap) < len(min_heap):
                max_heap.push(-min_heap.pop())
        median = -max_heap.top()
        if i >= K - 1:
            median_sum += median

    print(median_sum)


if __name__ == '__main__':
    main()

• Dependency: teflib.priorityqueue.UpdatableHeap