리트코드 543번 Diameter of Binary Tree (Python)

Problem

https://leetcode.com/problems/diameter-of-binary-tree/description/

Solution

1차 풀이

# Definition for a binary tree node.
from typing import Optional

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

class Solution:
    def diameterOfBinaryTree(self, root: Optional[TreeNode]) -> int:
        def search(node, depth):
            if node.left is None and node.right is None:
                return depth

            left_depth, right_depth = 0, 0
            if node.left is not None:
                left_depth = search(node.left, depth + 1)

            if node.right is not None:
                right_depth = search(node.right, depth + 1)

            return max(left_depth, right_depth)

        if root is None:
            return 0

        left_max_depth = right_max_depth = 0

        if root.left is not None:
            left_max_depth = search(root.left, 1)

        if root.right is not None:
            right_max_depth = search(root.right, 1)

        return left_max_depth + right_max_depth

root기준 왼쪽의 최대 깊이 + 오른쪽의 최대 깊이를 하면 된다고 생각했다.

하지만 위와 같이 경로에 루트가 포함되지 않는 경우에 대해 예외 케이스가 생긴다.

2차 풀이

class Solution:
    max_length = 0
    

    def diameterOfBinaryTree(self, root: Optional[TreeNode]) -> int:
        def dfs(node: TreeNode):
            if node.left is None and node.right is None:
                return 1

            left_depth, right_depth = 0, 0

            if node.left is not None:
                left_depth = dfs(node.left)
            if node.right is not None:
                right_depth = dfs(node.right)
                
            self.max_length = max(self.max_length, left_depth + right_depth)
            
            return max(left_depth, right_depth) + 1
        
        dfs(root)
        
        return self.max_length

leaf node까지 탐색한 후 위로 올라오면서 상태 값을 업데이트 해주었다.

노드의 높이를 전달하고, 현재 노드 기준으로 만들 수 있는 경로의 최대 길이를 계속해서 업데이트 했다.

Reference

파이썬 알고리즘 인터뷰 43번