优化 Python 中二叉树的查找直径

Optimize finding diameter of binary tree in Python

我想知道如何以最佳方式找到二叉树的直径(或任意两个叶节点之间的最长路径)。我有下面的基本解决方案,但第二种解决方案需要传递指针。我怎样才能在 Python 中做这样的事情?

def find_tree_diameter(node):
    if node == None:
        return 0

    lheight = height(node.left)
    rheight = height(node.right)

    ldiameter = find_tree_diameter(node.left)
    rdiameter = find_tree_diameter(node.right)

    return max(lheight+rheight+1, ldiameter, rdiameter)

def find_tree_diameter_optimized(node, height):
    lheight, rheight, ldiameter, rdiameter = 0, 0, 0, 0

    if node == None:
#        *height = 0;
        return 0

    ldiameter = diameterOpt(root.left, &lheight)
    rdiameter = diameterOpt(root.right, &rheight)

#    *height = max(lheight, rheight) + 1;

    return max(lh + rh + 1, max(ldiameter, rdiameter));

Python 支持多个 return 值,因此您不需要像 C 或 C++ 中那样的指针参数。这是代码的翻译:

def diameter_height(node):
    if node is None:
        return 0, 0
    ld, lh = diameter_height(node.left)
    rd, rh = diameter_height(node.right)
    return max(lh + rh + 1, ld, rd), 1 + max(lh, rh)

def find_tree_diameter(node):
    d, _ = diameter_height(node)
    return d

函数diameter_height return是树的直径和高度,find_tree_diameter用它来计算直径(通过丢弃高度)。

无论树的形状如何,函数都是O(n)。原函数在最坏的情况下是O(n^2),因为重复计算高度导致树非常不平衡。

简单Python 3 解

def findDepth(root):

    if root is None:
        return 0

    return 1 + max(findDepth(root.left), findDepth(root.right))

class Solution:
    def diameterOfBinaryTree(self, root: TreeNode) -> int:

        if root is None:
            return 0

        left = findDepth(root.left)
        right = findDepth(root.right)

        ldia = self.diameterOfBinaryTree(root.left)
        rdia = self.diameterOfBinaryTree(root.right)


        return max(left+right, max(ldia, rdia))