[백준 1572번, 9426번] 중앙값 측정 (Python3)

import sys
input = sys.stdin.readline
from heapq import heappush,heappop
INF = 10**6

N,K = map(int,input().split())

maxhq = []
minhq = []
heappush(maxhq,(INF,INF))
heappush(minhq,(INF,INF))

value = [0]*N
result = 0
for i in range(N):
  n = int(input())
  value[i] = n
  if n>minhq[0][0]:
    heappush(minhq,(n,i))
  else:
    heappush(maxhq,(-n,i))
  if i>=K:
    if maxhq[0][1] == i-K:
      heappop(maxhq)
    elif minhq[0][1] == i-K:
      heappop(minhq)
    else:
      if value[i-K]>minhq[0][0]:
        heappush(maxhq,(INF,INF))
      else:
        heappush(minhq,(INF,INF))
  while True:
    if maxhq[0][1]<=i-K:
      heappop(maxhq)
      heappush(maxhq,(INF,INF))
      continue
    if minhq[0][1]<=i-K:
      heappop(minhq)
      heappush(minhq,(INF,INF))
      continue
    if len(minhq)>len(maxhq):
      v,j = heappop(minhq)
      heappush(maxhq,(-v,j))
      continue
    if len(maxhq)>len(minhq)+1:
      v,j = heappop(maxhq)
      heappush(minhq,(-v,j))
      continue
    if i>=K-1:
      result -= maxhq[0][0]
    break
print(result)

가운데를 말해요 [백준 1655번] 가운데를 말해요 (Python3) (tistory.com) 문제와 거의 동일한 방식으로 풀었다.

최소힙, 최대힙을 사용하여 가운데값을 찾는 방식이다.

 

하지만 가운데를 말해요 문제와는 다르게 모든 범위에서의 중앙값이 아닌 특정 구간에서의 중앙값이기 때문에 약간의 아이디어가 필요하다.

 

내가 생각한 아이디어는 구간 밖의 값을 제거하지 못한다면 더미데이터를 추가하는 방식이었다.

 

간략하게 설명하면,

1. 기본적인 방식은 가운데를 말해요와 같다. 반 잘랐을때 큰 값을 최소힙에, 작은 값을 최대힙에 넣고, 입력받은 n값에 대해 최소힙의 가장 작은값보다 클 경우 최소힙에, 아닐 경우 최대힙에 음수로 넣는다.

2. 이때 그냥 값만을 넣는 것이 아니라 값과 인덱스로 이루어진 튜플로 삽입한다.

3. 구간 K 범위 밖을 벗어나는 값인 i-K번째 값을 제거해주어야 한다. 만약 minhq[0]나 maxhq[0]에 i-K번째 값이 있다면 pop하여 제거한다.

4. 제거하지 못했다면 해당 값이 최소힙에 들어있는지, 최대힙에 들어있는지를 파악하여 (minhq[0][0]와의 크기비교로 확인가능) 반대쪽 힙큐에 (INF,INF) 튜플을 더미데이터로 삽입한다.

5. while문을 돌려 가운데를 맞추어 준다. 이때 맞추는 과정중에 minhq[0]나 maxhq[0]에 범위 K 바깥의 값이 있다면 pop 해준 뒤 해당 힙에 (INF,INF) 튜플을 더미데이터로 삽입한다.

6. 가운데 값을 합하여 결과 출력

 

다른 사람들 풀이를 보니 세그먼트 트리를 이용한 풀이가 많아보였다. 힙큐를 이용한 풀이가 간단하지만 이후에는 또 어떤 문제가 나올지 모르니 해당 풀이도 알아둬야할 필요가 있을것이다.

 

 

오늘의 교훈) 세그먼트 트리를 이용한 중앙값 출력을 알아보자