[Leetcode] 리트코드 - 이진 트리의 직경(diameter-of-binary-tree) 파이썬(python) 풀이

문제

이진 트리에서 두 노드 간 가장 긴 경로의 길이를 출력하라. 다음과 같은 이진 트리가 주어졌을 때,

      1
     2 3
    4 5

가장 긴 경로는 4 → 2 → 1 → 3 또는 5 → 2 → 1 → 3으로 3이다.

코드

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


class Solution:
    longest: int = 0

    def diameterOfBinaryTree(self, root: TreeNode) -> int:
        def dfs(node: TreeNode) -> int:

            # 리프노드의 left, right 노드를 None으로 초기화하므로 거기까지 타고 갔을때 -1을 리턴됨.
            if not node:
                return -1

            # 리프노트 다음 가상노드까지 끝까지 타고 들어감
            left = dfs(node.left)
            right = dfs(node.right)

            # 왼쪽, 오른쪽 상태값에 2를 더해서 거리를 구함. 기존 거리에 누적으로 max 값을 측정.
            self.longest = max(self.longest, left + right + 2)

            # 아들 노드에 1더해서 상태값 리턴.
            return max(left, right) + 1
        dfs(root)

        return self.longest

풀이

가장 긴 경로를 찾는 방법은 먼저 가장 밑단, 즉 리프 노드까지 탐색한 다음 다음 부모로 거슬러 올라가면서 각각의 거리를 계산해 상태값을 업데이트하면서 누적해 나간다.

이렇게 거슬러 올라가 최종 루트에서 상태값은 2, 거리는 3이된다. 정답은 거리인 3이다. 상태값은 리프 노드에서 현재 노드까지의 거리다. 거리는 왼쪽 자식 노드의 상태값과 오른쪽 자식 노드의 상태값의 합에 2를 더한 값이다. 다시 정리하면, 최종적으로 거리는 왼쪽 자식 노드의 리프 노드에서 현재 노드까지의 거리와 , 오른쪽 자식 노드의 리프 노드에서 현재 노드까지의 거리의 합에 2를 더한 값과 같다.

def dfs(node: TreeNode) -> int:
    ...
    left = dfs(node.left)
    right = dfs(node.right)

이처럼 계속 재귀 호출을 통해 왼쪽, 오른쪽의 각 리프 노드까지 DFS로 탐색한다.

def dfs(node: TreeNode) -> int:
    ...
    self.longest = max(self.longest, left + right + 2)
    return max(left, right) + 1

이후에는 2개의 값을 계산하는데 ,하나는 최종 결과가 될 가장 긴 경로 self.longest, 나머지 하나는 앞서 얘기한 상태값 max(left, right) + 1을 말한다.

자식 노드가 하나도 없는 경우 left, right는 모두 -1이고, 이 경우 거리는 0 상태값도 0이 된다. 자식 노드가 모두 존재하는 경우에는, 그리고 자식 노드가 둘 다 상태값이 0이라면, 거리는 2, 상태값은 1이 된다. 즉 거리는 왼쪽, 오른쪽 자식 사이의 경로이므로 2를 더하게 되고, 상태값은 양쪽 자식 중 최대 상태값과 부모까지의 거리인 1을 더하게 된다. 그래서 예제의 입력값에서 루트의 상태값은 2, 거리값은 3이 된다. 정답은 거리값인 3이다.

예제의 노드에서 함수의 순서대로 따라가게 되면 각 변수들이 밑의 표처럼 나타나게 된다.

출처
파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]