[자료구조] [heap] 업데이트 되는 array값의 중앙값 구하기

Hackerrank의 Data structures 관한 문제 중에서 Hard에 속하는 문제이다. 원 문제 링크는 다음을 따라가보자. 

https://www.hackerrank.com/challenges/find-the-running-median/problem?h_r=internal-search

Find the Running Median | HackerRank

   연속적으로 주어지는 배열의 평균값들을 반복적으로 구하는 문제이다. [12,4,5,3,8,7]이 주어지면 앞에서부터 한 개씩 차례대로 추가하여 추가된 값들까지의 중앙값을 구하고, 이들이 포함된 리스트를 구하면 된다. 구체적인 연산과정을 살펴보자. 

1) [12] -> median: 12

2) [12,4] -> median: 8

3) [12,4,5] -> median : 5

4) [12,4,5,3] -> median: 4.5

와 같은 식으로 계산되어 정답값은 [12,8,5,4.5, 5, 6]이 나오도록 알고리즘을 짜야 한다. 하나씩 값이 추가될 때마다 1) 크기에 따른 정렬이 필요하며, 2) 정렬된 리스트의 길이가 홀수이면 리스트의 중앙값에 위치한 값을 그대로 호출하고, 리스트의 길이가 짝수라면 중간에 위치한 값들의 평균값을 호출해야 한다. 

먼저 값이 새로 추가될 때마다 정렬할 때 효율적으로 정렬하는 알고리즘은 없을까? 파이썬의 내장되어 있는 list.sort() 함수의 경우 time complexity가 O(N log N)이므로, 이보다 효율적인 자료구조형 혹은 알고리즘을 찾아야 한다. 

리스트에서 가장 작거나 큰 요소에 접근하는데 있어서 효율적인 데이터 타입인 heap을 생각해볼 수 있다. 다만, 위의 문제에서는 중앙값을 구하는 문제이기 때문에 찾고자 하는 중앙값을 기준으로 작은 값들을 포함하고 있는 heap 부분과 큰 값을 포함하고 있는 heap부분으로 나눠서 생각해야 한다. 

[12,4,5,3,8,7]일 경우 가능한 두 개의 힙

기본적으로, 파이썬이 제공하는 heapq 모듈을 이용하되, heapq.heappop(list)가 heap에서 가장 작은 항목을 제거 후 반환한다는 걸 염두하였다. 

Median을 기준으로 작은 값들을 포함하는 heap에서는 최대값을 뽑아야 하므로, heap에 넣을 때 각각의 값들에 -1을 넣어서 최상위 노드에 포함되는 값의 실제값이 가장 크도록 만든다. 입력할 때 -1을 넣었으므로 뽑아낼 때에는 다시 -1을 곱해 원래 값을 복원한다. 

import heapq

lowerThanMedian = []
higherThanMedian = []
median = 0

def putInMaxHeap(val):
    heapq.heappush(lowerThanMedian, -1*val)
    

def getFromMaxHeap():
    val = heapq.heappop(lowerThanMedian)
    return -1*val

또한, 하나씩 값을 추가할 때마다 각각의 lowerThanMedian과 higherThanMedian의 길이가 같아야 한다. 

if abs(len(lowerThanMedian) - len(higherThanMedian)) > 1:
  if len(lowerThanMedian) > len(higherThanMedian):
  	heapq.heappush(higherThanMedian, getFromMaxHeap())
  
  else:
  	putInMaxHeap(heapq.heappop(higherThanMedian))

 

마지막, median값을 출력하는데 있어서, 두 개의 heap을 합한 길이가 짝수라면 두 heap에서 뽑히는 최대값과 최소값을 2로 나누면 되고

홀수인 경우에는 두 개의 힙 중에서 길이가 큰 heap의 값을 호출하면 된다. heappop되는 값과 heap의 첫번째 값은 같다. 인덱스 0에 위치한 값이 최상위 노드로서 가장 작은 값과 큰 값을 가지고 있기 때문이다. 

        if (len(lowerThanMedian) + len(higherThanMedian)) % 2 == 0:
            median = (lowerThanMedian[0]*-1 + higherThanMedian[0])/2

        else:
            if len(lowerThanMedian) > len(higherThanMedian):
                median = lowerThanMedian[0]*-1
            else:
                median = higherThanMedian[0]

 

* 다른 식의 풀이

새로운 값을 추가할 때 이진 탐색을 기반으로 추가할 위치를 알려주는 python의 기본 모듈이 있다. bisect의 모듈에서 bisect라는 함수는 정렬된 리스트에 새로운 값을 추가할 index를 알려준다. 하지만, 이 알고리즘 역시 추후에는 list.insert(index, value)을 사용해야 한다. 이는 O(N)의 시간복잡도가 걸리기 때문에 list의 길이가 매우 길어질 경우, 별로 위의 방법보다 효율적이지 못하다.