[백준 / python] 13325: 이진트리

https://www.acmicpc.net/problem/13325

문제

각 에지에 양수인 가중치가 부여된 높이가 k인 포화이진트리가 주어져 있다. 높이 k인 포화이진트리는 2k개의 리프를 포함하여 (2k+1 − 1)개의 노드를 가진다. 루트에서 어떤 리프까지의 거리는 루트에서 그 리프까지의 경로상에 있는 모든 에지들의 가중치를 더한 값이다. 이 문제에서는, 어떤 에지들의 가중치를 증가시켜서 루트에서 모든 리프까지의 거리가 같도록 하고, 또한 에지 가중치들의 총합을 최소화 하려고 한다.

포화이진트리에 들어 있는 모든 에지들의 가중치가 주어졌을 때, 어떤 에지들의 가중치를 증가시켜서 루트에서 모든 리프까지의 거리가 같게 하면서 에지 가중치들의 총합이 최소가 되도록 하는 프로그램을 작성하시오.

입력

입력 데이터는 표준입력을 사용한다. 입력의 첫째 줄에는 포화이진트리의 높이를 나타내는 양의 정수 k(1 ≤ k ≤ 20)가 주어진다. 두 번째 줄에는 모든 에지들의 가중치가 주어진다. 에지들의 가중치는 루트에서 가까운 레벨에 있는 것부터, 같은 레벨에 있는 경우는 왼쪽부터 오른쪽의 순서로 주어진다. 각 에지의 가중치는 1 이상 1,000 이하인 정수로 주어진다.

출력

출력은 표준출력을 사용한다. 에지들의 가중치를 증가시킨 다음에 얻어지는 트리에 있는 모든 에지들의 가중치들의 총합을 한 줄에 출력한다. 어떤 에지의 가중치는 경우에 따라 1,000 이상의 값으로 증가될 수도 있다는 점에 유의하시오.

풀이

트리의 하단에서부터 양 옆의 간선의 가중치 중 가장 큰 가중치에 맞춰서 값을 바꿔주며 올라오면 된다.

import sys

k = int(sys.stdin.readline())
tree = [0] + list(map(int, sys.stdin.readline().strip().split()))

def dfs(index):
    left_index = index * 2 + 1
    right_index = index * 2 + 2

    if left_index >= len(tree):
        return 0

    left_edge, right_edge = tree[left_index], tree[right_index]
    left_max, right_max = dfs(left_index), dfs(right_index)
    left_result, right_result = left_edge + left_max, right_edge + right_max

    if left_result > right_result:
        tree[right_index] = left_result - right_max
    else:
        tree[left_index] = right_result - left_max

    return max(left_result, right_result)

dfs(0)
print(sum(tree))