Tree Cutting

• 트리에서 문제에서 주어진 성질을 만족하는 노드가 센트로이드이다.
• 링크에서 설명한대로 센트로이드를 찾으면 된다. 총 시간복잡도는 O(n).
• 센트로이드가 존재하지 않는 경우는 없으므로 'NONE'을 출력할 일은 없다. 그 대신, 센트로이드가 2개 존재하는 경우는 처리가 필요하다. 센트로이드 분할을 위해서는 센트로이드가 2개 있어도 아무거나 하나만 찾으면 되지만, 이 문제에서는 두개를 모두 찾아야 한다. 그리고 실수로 빼먹기 쉽지만, 두 센트로이드를 오름차순으로 출력해야 한다.
• 센트로이드가 2개 존재하는 경우는 전체 노드의 갯수 N이 짝수이고, 사이즈가 N/2 인 서브트리가 존재할때이다. 이때는 그 서브트리의 루트 노드와, 서브트리의 부모노드가 모두 센트로이드가 된다.
• 각 서브트리의 크기를 구하기 위해 teflib.tree.dp_on_tree 을 사용하면 한줄로 처리가 가능하긴 하지만, 그냥 최적화된 dfs를 돌리는것보다 속도는 느리다. 실행 시간이 112ms vs 168ms 로 차이가 났다.

"""Solution code for "BOJ 7045. Tree Cutting".

- Problem link: https://www.acmicpc.net/problem/7045
- Solution link: http://www.teferi.net/ps/problems/boj/7045

Tags: [Centroid]
"""

import sys
from teflib import tree as ttree


def find_centroids(tree):
    sizes = ttree.dp_on_tree(tree, 0, lambda vals, u: sum(vals) + 1)
    half_size = len(tree) // 2
    p, u = -1, 0
    while True:
        c = max((c for c in tree[u] if c != p), key=sizes.__getitem__)
        if sizes[c] == half_size:
            return [u, c]
        elif sizes[c] < half_size:
            return [u]
        p, u = u, c


def main():
    N = int(sys.stdin.readline())
    if N == 1:
        print('1')
        return
    tree = [[] for _ in range(N)]
    for _ in range(N - 1):
        X, Y = [int(x) for x in sys.stdin.readline().split()]
        tree[X - 1].append(Y - 1)
        tree[Y - 1].append(X - 1)
    print(*(u + 1 for u in sorted(find_centroids(tree))), sep='\n')


if __name__ == '__main__':
    main()

• Dependency: teflib.tree.dp_on_tree