파이썬으로 이진 트리의 대각선 경로 요소 합 구하기

문제 개요

이진 트리가 주어졌을 때, 트리의 왼쪽 위에서 오른쪽 아래 방향으로 내려가는 각 대각선 경로에 포함된 노드 값들의 합을 구하는 문제입니다.

예를 들어 다음과 같은 이진 트리가 있다고 가정해 보겠습니다.

이 트리의 대각선 경로는 [12, 15], [8, 10], [3] 세 가지이며, 따라서 출력 결과는 [27, 18, 3]이 됩니다.

접근 방법

핵심 아이디어는 재귀적으로 트리를 순회하면서, 왼쪽 자식으로 이동할 때는 대각선 번호를 1 증가시키고, 오른쪽 자식으로 이동할 때는 대각선 번호를 그대로 유지하는 것입니다.

알고리즘 단계

traverse() 함수를 정의합니다. 이 함수는 노드(node), 대각선 번호(numLeft), 결과 리스트(output)를 매개변수로 받습니다.

노드가 null이면 즉시 반환합니다.
numLeft가 output 리스트의 크기보다 크거나 같으면, 현재 노드의 데이터를 output의 끝에 새로 추가합니다.
그렇지 않으면, 기존 항목에 노드의 데이터를 더합니다. 즉, output[numLeft] += node.data
노드의 왼쪽 자식이 존재하면, traverse(왼쪽 자식, numLeft + 1, output)을 호출합니다.
노드의 오른쪽 자식이 존재하면, traverse(오른쪽 자식, numLeft, output)을 호출합니다.

메인 메서드에서는 다음을 수행합니다.

빈 리스트 output을 생성합니다.
traverse(root, 0, output)을 호출하여 순회를 시작합니다.
output을 반환합니다.

구현 코드

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

class Solution:
    def solve(self, root):
        output = []

        def traverse(node, numLeft, output):
            if not node:
                return
            if numLeft >= len(output):
                output.append(node.data)
            else:
                output[numLeft] += node.data
            if node.left:
                traverse(node.left, numLeft + 1, output)
            if node.right:
                traverse(node.right, numLeft, output)

        traverse(root, 0, output)
        return output

ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
print(ob.solve(root))

입력

root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)

출력

[27, 18, 3]

복잡도 분석

이 알고리즘은 트리의 모든 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(N)입니다. 여기서 N은 트리의 노드 수입니다. 공간 복잡도는 재귀 호출 스택 깊이에 비례하며, 최악의 경우 편향된 트리일 때 O(N), 균형 잡힌 트리일 때 O(log N)입니다.