[HackerRank] Find the Running Median (Python)

1. 문제

https://www.hackerrank.com/challenges/find-the-running-median/problem

 

2. 풀이

처음엔 a의 길이가 홀수, 짝수 케이스로 나누어서 단순히 중위값을 구하려고했다.

참고로 python에서 n / 2 를 하면 실수를 반환하지만 n//2를 하면 int 정수를 반환한다.

def runningMedian(a):
    answer = []
 
    for i in range(1, len(a)+1):
        li = sorted(a[:i])
        n = len(li)
        
        # odd
        if (n % 2 != 0):
            median = li[n // 2]
        else:
            median = (li[n // 2] + li[(n // 2) - 1]) / 2

        answer.append(median)

    return answer

 

하지만 위 코드의 시간복잡도는 

정렬(sorted)은 O(k log k), 여기서 k는 배열의 길이, 총 n번 반복하므로 시간 복잡도는 O(n² log n) 이다.

 

heap 을 이용해야 시간복잡도를 줄일 수 있는 문제. Heap을 사용하면 시간 복잡도를 O(nlogn)으로 줄일 수 있다.

- left heap : 중위값보다 작은 값을 저장

- right heap : 중위값보다 큰 값을 저장

heap 의 insert , delete 가 O(logN) 이므로 n 번 반복하면 O(nlogn) 이다.

 

중위값을 구할 때, 왼쪽 힙이 더 크거나 두 힙의 크기가 같으면 왼쪽 힙의 루트가 중위값이고, 그렇지 않으면 두 힙의 루트를 평균하여 중위값을 구합니다.

2개의 heap 을 사용하므로 공간 복잡도는 O(N) 이다.

 

import heapq 필요!!!

 

3.코드

import heapq

def runningMedian(a):
    # 두 개의 힙을 사용할 것: 
    # max heap (left)와 min heap (right)
    left = []  # max heap (파이썬은 기본적으로 min heap이라 음수로 값 저장)
    right = []  # min heap
    
    result = []
    
    for num in a:
        # 새로운 숫자가 들어오면 두 힙 중 적절한 곳에 추가
        if len(left) == 0 or num <= -left[0]:
            heapq.heappush(left, -num)  # max heap에 삽입 (음수로 저장)
        else:
            heapq.heappush(right, num)  # min heap에 삽입
        
        # 두 힙의 크기 균형 맞추기
        if len(left) > len(right) + 1:
            # 왼쪽 힙이 너무 크면, 왼쪽에서 하나를 오른쪽으로 이동
            heapq.heappush(right, -heapq.heappop(left))
        elif len(right) > len(left):
            # 오른쪽 힙이 너무 크면, 오른쪽에서 하나를 왼쪽으로 이동
            heapq.heappush(left, -heapq.heappop(right))
        
        # 중위값 계산
        if len(left) > len(right):
            # 왼쪽 힙이 더 크면, 중위값은 왼쪽 힙의 루트
            result.append(float(-left[0]))
        else:
            # 두 힙의 크기가 같으면, 중위값은 두 힙의 루트의 평균
            result.append((-left[0] + right[0]) / 2.0)
    
    return result

 

• left (max heap):
  • heapq는 기본적으로 min heap을 사용하므로, 중위값보다 작은 값들은 음수로 저장하여 max heap처럼 동작하도록 합니다.
• right (min heap):
  • heapq에서 기본적으로 제공하는 min heap을 그대로 사용하여 중위값보다 큰 값들을 저장합니다.
• 두 힙의 크기 균형:
  • 두 힙의 크기가 너무 차이나지 않도록, 힙이 한쪽으로 치우칠 경우 크기가 큰 힙에서 하나를 작은 힙으로 이동시킵니다.
• 중위값 계산:
  • left 힙의 크기가 right 힙보다 크면 left 힙의 루트가 중위값입니다.
  • 두 힙의 크기가 같으면 두 힙의 루트 값을 평균으로 계산합니다.