중앙값 측정

• 데이터가 삽입, 삭제될때마다 median을 갱신하는 문제
• Running Median 방법으로 풀 수 있는 문제이다. 자세한 설명은 링크 참조.
• 가운데를 말해요와도 비슷하지만, 여기에서는 삽입 외에 삭제 처리도 같이 해주어야 한다는 점이 다르다.
• 중앙값와는 동일한 문제.
• 힙에서 삭제 연산을 처리하기 위해서 teflib.priorityqueue.UpdatableHeap 을 사용했는데, 이것은 실제로 원소를 삭제하는 대신 원소에 마킹을 해놓고 나중에 원소가 top에 올라오면 삭제하는 구조이기 때문에, 이 문제의 경우 실제로 힙에 들어가있을수 있는 원소의 갯수는 K가 아닌 n까지 가능하다. 따라서 삽입 연산의 시간 복잡도가 O(logK)가 아닌 O(logn)이 되지만 큰 차이는 아니다.
• n번의 삽입, 삭제, 중앙값 쿼리를 처리하는데에 걸리는 총 시간 복잡도는 O(nlogn)
• 삭제를 처리하기 위해서 추가 작업들이 많이 있으므로, Order Statistic Tree를 써도 많이 안느려지지 않을까 생각해봤지만, 시도한 결과는 힙을 썼을때 걸린 1060ms 보다 훨씬 느린 2460ms 정도에 통과되었다.
  • 오히려 다른 사람들의 코드에서, 그냥 리스트에 K개의 값을 정렬시켜서 저장하면서 O(K)에 삽입과 삭제를 처리하는 방법이, 시간복잡도는 안좋지만 실행시간은 더 빠른 900ms대에 통과되는 것을 발견했다.

"""Solution code for "BOJ 9426. 중앙값 측정".

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

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 max_heap.size() > min_heap.size() + 1:
                min_heap.push(-max_heap.pop())
        else:
            min_heap.push(num)
            while max_heap.size() < min_heap.size():
                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