如何找到二叉树的直径？

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

def find_diameter(root):
    if root is None:
        return 0

    max_diameter = 0

    def dfs(node):
        nonlocal max_diameter
        if node is None:
            return 0

        left_height = dfs(node.left)
        right_height = dfs(node.right)

        max_diameter = max(max_diameter, left_height + right_height)

        return max(left_height, right_height) + 1

    dfs(root)

    return max_diameter