0543__diameter_of_binary_tree.py

"""
Given the root of a binary tree, return the length of the diameter of the tree.

The diameter of a binary tree is the length of the longest path between any two nodes in a tree. This path may or may
not pass through the root.

The length of a path between two nodes is represented by the number of edges between them.

## Example 1

Input: root = [1,2,3,4,5]
Output: 3
Explanation: 3 is the length of the path [4,2,1,3] or [5,2,1,3].

## Example 2

Input: root = [1,2]
Output: 1

## Constraints

* The number of nodes in the tree is in the range [1, 10^4].
* -100 <= Node.val <= 100
"""
from typing import Optional
from unittest import TestCase

from lib.TreeNode import TreeNode, build_tree


class Solution(TestCase):
    max: int = 0

    def test_example_1(self):
        self.assertEqual(3, self.diameterOfBinaryTree(build_tree([1, 2, 3, 4, 5])))

    def test_example_2(self):
        self.assertEqual(1, self.diameterOfBinaryTree(build_tree([1, 2, None])))

    def test_empty_tree(self):
        self.assertEqual(0, self.diameterOfBinaryTree(None))

    def test_only_root(self):
        self.assertEqual(0, self.diameterOfBinaryTree(build_tree([1])))

    def diameterOfBinaryTree(self, root: Optional[TreeNode]) -> int:
        self.max = 0

        def dfs(node: Optional[TreeNode]) -> int:
            if not node:
                return 0

            left, right = dfs(node.left), dfs(node.right)
            path = left + right

            self.max = max(self.max, path)

            return 1 + max(left, right)

        dfs(root)

        return self.max