[python] algospot RUNNINGMEDIAN

문제

-> 문제 링크 : https://algospot.com/judge/problem/read/RUNNINGMEDIAN

algospot.com :: RUNNINGMEDIAN

 

문제풀이

수열의 길이가 최대 20만이기 때문에 매번 새로운 숫자가 들어왔을 때 중간값을 구하는 방법으로는 시간 안에 해결할 수 없다.

O(nlogn)이나 O(logn)에 해결할 수 있는 방법이 필요하다.

수열에 있는 숫자들은 모두 한 번씩은 사용해야 하기 때문에 n이 나올 수밖에 없다고 생각했고, 숫자 하나를 삽입하는 과정에서 logn의 시간으로 해결해야 한다고 생각했다.

 

두 개의 힙을 사용하여 minHeap과 maxHeap을 만든다. maxHeap은 항상 minHeap에 있는 값들보다 작은 값들 만 존재하고,

숫자가 계속해서 들어올 때 두 힙의 크기를 동일하게 맞춰준다면  두 힙의 루트에 있는 값 중에 항상 median이 존재하게 된다.

문제에서 수열의 크기가 짝수이면 두 수중에 작은 값이 중간값이 된다고 했다. 항상 새로운 값이 들어올 때마다 더 maxHeap의 루트가 minHeap의 루트보다 작은 값이 오게 한다면 maxHeap의 루트 값은 항상 중간값이 된다.

힙의 삽입과 삭제는 O(logn)의 시간이 소요되므로 제한 시간 안에 문제를 해결할 수 있다.

 

 

코드

import heapq
import sys

T = int(sys.stdin.readline().rstrip())

def input_generator(n,a,b) :

    yield 1983
    prev= 1983
    for _ in range(n-1) :
        now = (prev * a + b) % 20090711
        yield now
        prev = now

for _ in range(T) :

    N, a, b = map(int, sys.stdin.readline().rstrip().split(" "))

    minHeap = []
    minSize = 0
    maxHeap = []
    maxSize = 0
    inputs = input_generator(N,a,b)
    medians = []
    for _ in range(N) :

        num = next(inputs)

        if minSize == maxSize :
            heapq.heappush(maxHeap, -num)
            maxSize += 1
        else :
            heapq.heappush(minHeap, num)
            minSize += 1

        if minHeap and maxHeap :
            minTop = abs(minHeap[0])
            maxTop = abs(maxHeap[0])

            if minTop < maxTop :
                maxTop = -heapq.heappop(maxHeap)
                minTop = -heapq.heappop(minHeap)

                heapq.heappush(maxHeap, minTop)
                heapq.heappush(minHeap, maxTop)

        medians.append(-maxHeap[0] % 20090711)

    print(sum(medians) % 20090711)

 

초기 접근 방식과 오답의 원인

힙을 이런 방식으로 사용해 문제를 해결할 수 있다고 생각 못했다.

힙의 응용을 공부해봐야겠다.

 

코드에 잘못된 부분이나 개선 사항이 있다면 댓글로 알려주시면 감사하겠습니다.