가운데를 말해요 - 1655

1. 개요

문제 링크: https://www.acmicpc.net/problem/1655

2. 코드

import sys

def swap(a, b, heap_type):
    global left_heap, right_heap
    heap = []
    if heap_type == 'left':
        heap = left_heap
    elif heap_type == 'right':
        heap = right_heap
    temp = heap[a]
    heap[a] = heap[b]
    heap[b] = temp

def upHeapify(idx, heap_type):
    global left_size, right_size, left_heap, right_heap
    heap = []
    size = 0
    if heap_type == 'left':
        heap = left_heap
        size = left_size
    elif heap_type == 'right':
        heap = right_heap
        size = right_size

    parent = idx // 2
    if heap_type == 'left': #Max Heap
        while parent >= 1:
            if heap[parent] > heap[idx]:
                return
            swap(parent, idx, heap_type)
            idx = parent
            parent = parent // 2
    elif heap_type == 'right': #Min Heap
        while parent >= 1:
            if heap[parent] < heap[idx]:
                return
            swap(parent, idx, heap_type)
            idx = parent
            parent = parent // 2

def downHeapify(idx, heap_type):
    global left_size, right_size, left_heap, right_heap
    heap = []
    size = 0
    if heap_type == 'left':
        heap = left_heap
        size = left_size
    elif heap_type == 'right':
        heap = right_heap
        size = right_size

    left = idx * 2
    if heap_type == 'left': #Max Heap
        while left <= size:
            right = left + 1
            if right <= size:
                if heap[left] < heap[right]:
                    left = right
            if heap[left] < heap[idx]:
                return
            swap(left, idx, heap_type)
            idx = left
            left = left * 2
    elif heap_type == 'right': #Min Heap
        while left <= size:
            right = left + 1
            if right <= size:
                if heap[left] > heap[right]:
                    left = right
            if heap[left] > heap[idx]:
                return
            swap(left, idx, heap_type)
            idx = left
            left = left * 2

def addNode(data, heap_type):
    global left_size, right_size, left_heap, right_heap
    heap = []
    size = 0
    if heap_type == 'left':
        heap = left_heap
        left_size += 1
        size = left_size
    elif heap_type == 'right':
        heap = right_heap
        right_size += 1
        size = right_size

    heap.append(data)
    upHeapify(size, heap_type)

def pop(heap_type):
    global left_size, right_size, left_heap, right_heap
    heap = []
    size = 0
    if heap_type == 'left':
        heap = left_heap
        left_size -= 1
        size = left_size
    elif heap_type == 'right':
        heap = right_heap
        right_size -= 1
        size = right_size

    root = heap[1]
    heap[1] = heap[size + 1]
    downHeapify(1, heap_type)
    heap.pop()
    return root

left_size = 0
right_size = 0
n = int(input())
left_heap = [0]
right_heap = [0]
for i in range(0, n):
    x = int(sys.stdin.readline().rstrip())
    
    # initial conditional process
    if len(left_heap) == 1 and len(right_heap) == 1:
        addNode(x, 'left')
        print(left_heap[1])
        continue
    elif len(left_heap) == 2 and len(right_heap) == 1:
        if x < left_heap[1]:
            temp = pop('left')
            addNode(temp, 'right')
            addNode(x, 'left')
        else:
            addNode(x, 'right')
        print(left_heap[1])
        continue

    # conditional processes
    if len(left_heap) == len(right_heap):
        if right_heap[1] < x:
            temp = pop('right')
            addNode(temp, 'left')
            addNode(x, 'right')
        elif x < left_heap[1]:
            addNode(x, 'left')
        else:
            addNode(x, 'left')
    elif len(left_heap) > len(right_heap):
        if x < left_heap[1]:
            temp = pop('left')
            addNode(temp, 'right')
            addNode(x, 'left')
        elif x > right_heap[1]:
            addNode(x, 'right')
        else:
            addNode(x, 'right')
    elif len(left_heap) < len(right_heap):
        if x > right_heap[1]:
            temp = pop('right')
            addNode(temp, 'left')
            addNode(x, 'right')
        elif x < left_heap[1]:
            addNode(x, 'left')
        else:
            addNode(x, 'left')
    print(left_heap[1])

3. 풀이 메모

두 개의 우선순위힙을 활용한다. 한 힙은 최대힙(이하 left_heap), 다른 한 힙은 최소힙(이하 right_heap)이 되며, left_heap의 최상위 요소가 중간값이 되도록 조건 분기를 꼼꼼하게 작성한다.

조건 분기는 코드에서 직관적으로 확인할 수 있다. 따라서 이 곳에 수도코드를 작성하는 과정은 생략한다.

4. 코멘트

힌트로 '힙 두개를 사용해야 한다' 는 정보를 얻은 뒤 추가 정보 없이 문제 풀이를 진행했다. 알고리즘은 정확하게 작동하지만, 코드가 다소 난잡해진 것 같다. 가까운 시일 내에 heapq 패키지를 활용해 코드를 더 줄이고, 불필요한 조건 분기를 제거해 코드를 깔끔하게 바꾼 뒤 별도의 포스트를 작성하려고 한다.