[파이썬] 트리 순회의 응용 (LCA, 서브 트리 크기 등)

트리 순회(Tree Traversal)는 트리를 구성하는 노드들을 특정한 방식으로 순서대로 방문하는 것을 의미합니다. 전통적으로는 세 가지 종류의 트리 순회가 가장 널리 사용됩니다: 전위 순회(Preorder Traversal), 중위 순회(Inorder Traversal), 후위 순회(Postorder Traversal). 이러한 기본적인 트리 순회 방식 외에도, 트리 순회의 응용 기법을 사용하여 더 다양한 문제를 풀 수 있습니다.

이번 글에서는 트리 순회의 응용 기법 중 일부인 최소 공통 조상(Lowest Common Ancestor, LCA)서브 트리 크기(Subtree Size)에 대해 알아보겠습니다. 또한, 이러한 응용 기법을 Python 언어를 사용하여 구현하는 방법을 예시 코드와 함께 설명하겠습니다.

최소 공통 조상(Lowest Common Ancestor, LCA)

최소 공통 조상은 두 개의 노드의 가장 가까운 공통 조상을 찾는 문제입니다. 트리의 각 노드에 대해, 가장 가까운 공통 조상을 찾기 위해 두 개의 노드가 필요합니다. 최소 공통 조상은 여러 가지 방법으로 풀 수 있지만, 이 중 하나인 이진 리프트(LCA using Binary Lifting) 알고리즘을 살펴보겠습니다.

class LCA:
    def __init__(self, tree):
        self.tree = tree
        self.n = len(tree)
        self.jump = [[-1] * self.n for _ in range(int(math.log2(self.n)) + 1)]
        self.depth = [0] * self.n
        self.visited = [False] * self.n

    def dfs(self, node, d):
        self.visited[node] = True
        self.depth[node] = d

        for child in self.tree[node]:
            if not self.visited[child]:
                self.jump[0][child] = node
                self.dfs(child, d + 1)

    def build(self):
        self.dfs(0, 0)

        for j in range(1, int(math.log2(self.n)) + 1):
            for i in range(self.n):
                if self.jump[j - 1][i] != -1:
                    self.jump[j][i] = self.jump[j - 1][self.jump[j - 1][i]]

    def find_lca(self, a, b):
        if self.depth[a] > self.depth[b]:
            a, b = b, a

        for i in range(int(math.log2(self.n)), -1, -1):
            if self.depth[b] - self.depth[a] >= (1 << i):
                b = self.jump[i][b]

        if a == b:
            return a
        
        for i in range(int(math.log2(self.n)), -1, -1):
            if self.jump[i][a] != self.jump[i][b]:
                a = self.jump[i][a]
                b = self.jump[i][b]

        return self.jump[0][a]

위의 예시 코드는 이진 리프트 알고리즘을 사용하여 LCA를 찾는 방법을 구현한 것입니다. 방문하지 않은 노드를 DFS로 방문하며, 각 노드의 깊이와 부모 노드 정보를 저장합니다. 그리고 이 정보를 바탕으로 LCA를 계산합니다.

서브 트리 크기(Subtree Size)

서브 트리 크기는 각 노드를 루트로 하는 서브 트리의 크기를 구하는 문제입니다. 서브 트리 크기를 계산하려면, 각 노드에 대해 DFS를 사용하여 자식 노드의 개수를 세는 방식으로 구현할 수 있습니다.

class SubtreeSize:
    def __init__(self, tree):
        self.tree = tree
        self.n = len(tree)
        self.subtree_size = [0] * self.n
        self.visited = [False] * self.n

    def dfs(self, node):
        self.visited[node] = True
        size = 1

        for child in self.tree[node]:
            if not self.visited[child]:
                size += self.dfs(child)

        self.subtree_size[node] = size
        return size

    def calculate(self):
        self.dfs(0)

위의 예시 코드는 DFS를 사용하여 각 노드를 루트로 하는 서브 트리의 크기를 계산하는 방법을 구현한 것입니다. 한 번의 DFS로 트리를 순회하며, 자식 노드의 개수를 재귀적으로 더해 나갑니다.

마치며

위에서 소개한 예시 코드를 활용하면, 트리 순회의 응용 기법인 최소 공통 조상과 서브 트리 크기를 효율적으로 계산할 수 있습니다. 이러한 응용 기법은 트리 기반의 다양한 알고리즘에서 유용하게 사용될 수 있으며, 알고리즘 문제 해결에 도움이 될 수 있습니다. Python 언어를 사용하여 구현한 예시 코드를 참고하여, 실제 문제에 적용해 보시기 바랍니다.