본문 바로가기
개발/알고리즘

[백준 11279][Python] 최대힙 자료구조 직접 구현하기, 시간초과 나는 이유

by JeonJaewon 2021. 5. 29.

 입력 받을 때 input() 을 사용하면 시간초과가 날 수 있습니다. 저는 7%정도에서 시간초과 오류가 났는데, sys.stdin.readline()으로 바꾸니 한번에 통과했습니다.

 힙에 추가하는 부분보다는 삭제하는 부분이 어려운데, 자식 노드가 존재하는지 검사하지 않으면 list out of range 에러가 날 수 있기 때문에 size 변수를 잘 이용해서 현재 노드의 개수를 잘 파악해야 합니다. 힙은 완전이진트리이기 때문에 노드 개수만으로도 트리 모양을 예측할 수 있습니다.

 파이썬을 오랜만에 썼는데 까먹은 부분이 많아서 놀랐고 앞으로도 종종 써야겠다고 생각했습니다

import sys


class Heap:
    def __init__(self, max_size):
        self.size = 0
        self.data = [0 for i in range(0, max_size)]

    def push(self, item):
        self.size += 1
        self.data[self.size] = item;
        self.swap();

    def swap(self):
        index = self.size
        while index != 1:
            parentItem = self.data[int(index / 2)]
            currentItem = self.data[index]
            if parentItem < currentItem:
                tmp = currentItem
                self.data[index] = parentItem
                self.data[int(index / 2)] = tmp
                index = int(index / 2)
            else:
                break
        
    def delete(self):
        if self.size == 0:
            print(0);
            return

        print(self.data[1]);
        self.data[1] = self.data[self.size]
        self.data[self.size] = 0
        self.size -= 1
        parent = 1
        child = 2
        while child <= self.size:
            if (child + 1 <= self.size) and (self.data[child + 1] > self.data[child]):
                child = child + 1
            if self.data[child] > self.data[parent]:
                tmp = self.data[child]
                self.data[child] = self.data[parent]
                self.data[parent] = tmp
                parent = child
                child = parent * 2
            else:
                break

N = int(input())
heap = Heap(N);
for i in range(0, N):
    num = int(sys.stdin.readline())
    if num == 0:
        heap.delete()
    else:
        heap.push(num)

 

댓글