Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 이진 트리에서 합이 짝수인 가장 긴 경로의 길이 구하기


문제 소개

하나의 이진 트리(binary tree)가 주어졌을 때, 경로에 포함된 노드 값들의 합이 짝수가 되는 경로 중 가장 긴 것의 길이를 구하는 문제입니다.

예를 들어 루트가 2이고, 왼쪽 자식이 5, 오른쪽 자식이 4이며, 4의 왼쪽 자식은 8, 오른쪽 자식은 2, 그리고 8의 왼쪽 자식이 5인 트리가 주어졌다고 해보겠습니다. 이때 경로 [5, 2, 4, 8, 5]의 합은 24(짝수)이므로 정답은 5가 됩니다.

해결 전략: 깊이 우선 탐색(DFS)

이 문제는 DFS(깊이 우선 탐색)를 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 노드에서 다음 두 가지 값을 반환하는 것입니다.

  • 첫 번째 값: 해당 노드에서 시작해 아래 방향으로 뻗는 경로 중, 합이 짝수인 가장 긴 경로의 길이
  • 두 번째 값: 같은 조건에서 합이 홀수인 가장 긴 경로의 길이

노드 값이 홀수인지 짝수인지에 따라 자식에게서 받은 경로와 현재 노드를 연결했을 때 합의 홀짝성이 달라집니다. 즉, 홀수 값 노드는 홀수 합 경로를 짝수 합 경로로 만들고, 짝수 값 노드는 홀짝성을 그대로 유지합니다. 이 성질을 이용해 왼쪽과 오른쪽 자식의 경로를 교차하여 결합하면 조건을 만족하는 최장 경로를 찾을 수 있습니다.

알고리즘 단계

  1. dfs(node) 함수를 정의합니다.
  2. node가 null이면 (0, -inf) 쌍을 반환합니다. 빈 경로의 합은 0(짝수)이고, 홀수 합 경로는 존재하지 않기 때문입니다.
  3. (left_0, left_1) := dfs(왼쪽 자식), (right_0, right_1) := dfs(오른쪽 자식)을 호출합니다.
  4. 노드 값이 홀수라면:
    ans := max(ans, left_1 + right_0 + 1, left_0 + right_1 + 1)
    그리고 (max(left_1 + 1, right_1 + 1, 0), max(left_0 + 1, right_0 + 1))을 반환합니다.
  5. 노드 값이 짝수라면:
    ans := max(ans, left_0 + right_0 + 1, left_1 + right_1 + 1)
    그리고 (max(left_0 + 1, right_0 + 1, 0), max(left_1 + 1, right_1 + 1))을 반환합니다.
  6. 메인 로직에서는 ans를 0으로 초기화한 뒤 dfs(root)를 호출하고, 최종적으로 ans를 반환합니다.

여기서 ans는 지금까지 발견한 '합이 짝수인 가장 긴 경로'의 길이를 저장하는 변수입니다. 경로는 한 노드에서 꺾일 수 있으므로, 왼쪽 아래로 내려갔다가 현재 노드를 거쳐 오른쪽 아래로 내려가는 형태의 경로도 모두 고려해야 합니다.

파이썬 구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.val = data
        self.left = left
        self.right = right

class Solution:
    def solve(self, root):
        def dfs(node):
            if not node:
                return 0, float("-inf")
            left_0, left_1 = dfs(node.left)
            right_0, right_1 = dfs(node.right)
            if node.val & 1:
                self.ans = max(self.ans,
                               left_1 + right_0 + 1,
                               left_0 + right_1 + 1)
                return (max(left_1 + 1, right_1 + 1, 0),
                        max(left_0 + 1, right_0 + 1))
            else:
                self.ans = max(self.ans,
                               left_0 + right_0 + 1,
                               left_1 + right_1 + 1)
                return (max(left_0 + 1, right_0 + 1, 0),
                        max(left_1 + 1, right_1 + 1))

        self.ans = 0
        dfs(root)
        return self.ans

ob = Solution()
root = TreeNode(2)
root.left = TreeNode(5)
root.right = TreeNode(4)
root.right.left = TreeNode(8)
root.right.right = TreeNode(2)
root.right.left.left = TreeNode(5)
print(ob.solve(root))

입력

root = TreeNode(2)
root.left = TreeNode(5)
root.right = TreeNode(4)
root.right.left = TreeNode(8)
root.right.right = TreeNode(2)
root.right.left.left = TreeNode(5)

출력

5

복잡도 분석

모든 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(n)입니다. 재귀 호출 스택의 최대 깊이는 트리의 높이에 비례하므로 공간 복잡도는 O(h)입니다. 여기서 n은 노드의 개수, h는 트리의 높이를 의미합니다.