SWEA 5177 이진 힙 (python)

5177. [파이썬 S/W 문제해결 기본] 8일차 - 이진 힙 D2

문제

이진 최소힙은 다음과 같은 특징을 가진다.

- 항상 완전 이진 트리를 유지하기 위해 마지막 노드 뒤에 새 노드를 추가한다.

- 부모 노드의 값<자식 노드의 값을 유지한다. 새로 추가된 노드의 값이 조건에 맞지 않는 경우, 조건을 만족할 때까지 부모 노드와 값을 바꾼다.

- 노드 번호는 루트가 1번, 왼쪽에서 오른쪽으로, 더 이상 오른쪽이 없는 경우 다음 줄로 1씩 증가한다.

예를 들어 7, 2, 5, 3, 4, 6이 차례로 입력되면 다음과 같은 트리가 구성된다.

이때 마지막 노드인 6번의 조상은 3번과 1번 노드이다.

1000000이하인 N개의 서로 다른 자연수가 주어지면 입력 순서대로 이진 최소힙에 저장하고, 마지막 노드의 조상 노드에 저장된 정수의 합을 알아내는 프로그램을 작성하시오.

[입력]

첫 줄에 테스트케이스의 수 T가 주어진다. 1<=T<=50
다음 줄부터 테스트 케이스의 별로 N이 주어지고, 다음 줄에 1000000이하인 서로 다른 N개의 자연수가 주어진다. 5<=N<=500

[출력]

각 줄마다 "#T" (T는 테스트 케이스 번호)를 출력한 뒤, 답을 출력한다.

풀이과정

맨처음 접근은 먼저 정렬시킨 다음에 함수를 만들어서 순서대로 채울 수 있는 것으로 생각했다.

하지만 정렬이아닌 부모 노드보다 자식 노드가항상 작아야 하므로 정보를 계속해서 넘겨주는 방식으로 풀어야 한다는 것을 알게 되었다.

풀이

import sys
sys.stdin = open("input.txt")

T = int(input())

class Heap:
    def __init__(self):
        # 초기에 0을 집어넣어 idx를  1부터 시작할 수 있게 한다.
        self.arr = [0]
        # 0을 무시한 길이를 위해서 임의로 설정해 준다.
        self.length = 0

    def make_heap(self, num):
        # 들어온 숫자를 리스트에 넣어준다.
        self.arr.append(num)
        # 길이를 추가해 준다.
        self.length += 1

        # 현재 인덱스를 쉽게 사용하기 위해서 지정해 준다.
        idx = self.length

        # idx가 들어올 때 마다 검사하기 위해서
        # 마지막 idx부터 부모 노드를 찾아가는 과정이다.
        while idx > 1:
            # 부모 idx를 만들어주자
            pidx = idx // 2
            # 만약 부모노드가 더 크다면 자식 노드랑 바꿔주자
            if self.arr[pidx] > self.arr[idx]:
                self.arr[pidx], self.arr[idx] = self.arr[idx], self.arr[pidx]

            # 부모 단계로 거슬러 올라가 보자
            idx = pidx



    # 마지막 노드의 조상들의 합을 구해보자
    def last(self, num):

        # 마지막 노드번호는 입력값과 같다
        idx = num
        # 총합을 구할 변수 초기화
        total = 0
        # 부모 찾기위해 설정
        while idx > 1 :
            pidx = idx // 2
            total += self.arr[pidx]
            idx = pidx

        return total


for tc in range(1, T+1):
    N = int(input())
    numbers = list(map(int, input().split()))

    heap = Heap()

    for number in numbers:
        heap.make_heap(number)

    result = heap.last(N)
    print("#{} {}".format(tc, result))

추가

입력값 받을때 귀찮아서 input().split()으로만 받았다가 피봤다...

str형식으로 넘어가서 16같은 경우가 1이라는 숫자가 더 빠르기 때문에 정렬이 이상하게 됬다...