在二叉搜索树中递归统计三种类型节点的数量

Question

我正在尝试写一个函数，用来计算二叉搜索树中三种类型的节点数量，分别是没有孩子的节点、只有一个孩子的节点和有两个孩子的节点。经过观察，我发现用递归的方法来实现这个功能会比用循环的方法更好。

def node_counts(self):
    """
    ---------------------------------------------------------
    Returns the number of the three types of nodes in a BST.
    Use: zero, one, two = bst.node_counts()
    -------------------------------------------------------
    Postconditions:
      returns
      zero - number of nodes with zero children (int)
      one - number of nodes with one child (int)
      two - number of nodes with two children (int)
    ----------------------------------------------------------
    """
    zero, one, two = self._node_counts_aux(self._root)
    return zero, one, two

def _node_counts_aux(self, node):
    zero, one, two = 0, 0, 0
    if node is not None:
        if not node._right and not node._left:
            zero = 1 # I understand that the problem is here.
        if node._left and node._right:
            two = 1 + self._node_counts_aux(node._left)[2] + self._node_counts_aux(node._right)[2]
        if node._left or node._right:
            one = 1 + self._node_counts_aux(node._left)[1] + self._node_counts_aux(node._right)[1]
    return zero, one, two

"""
I am testing with this Tree:
        36
       /  \
      /    \
     6      50
    / \     / \
   4   17  49  84
       /   /   / \
      12  42  65 85

The output with this code comes to: (0, 6, 4).

"""

这里有一列数据在某种意义上是错的，但在另一种意义上又是对的，这不是我关心的重点。我的重点是零没有被计算在内。零被设置为0，那我该怎么解决这个问题呢？

Answer 1

你需要把从递归调用中得到的结果累加起来。可以用 zero, one, two = map(sum, zip(result_right, result_left)) 这个方法来实现，然后根据孩子的数量来加上相应的值。

注意，我使用了 if/elif 语句，否则当节点有两个孩子时，你的代码也会进入下一个只针对一个孩子的 if 块。

def _node_counts_aux(self, node):
    zero, one, two = 0, 0, 0
    if node is not None:
        result_right = self._node_counts_aux(node._right)
        result_left = self._node_counts_aux(node._left)
        zero, one, two = map(sum, zip(result_right, result_left))
        if not node._right and not node._left:
            zero += 1
        elif node._left and node._right:
            two += 1
        elif node._left or node._right:
            one += 1
    return zero, one, two

Answer 2

问题在于，方法 _node_counts_aux() 返回的是一个元组，但你却试图在它的结果上加 1。你需要从递归调用中提取出类型为 0、1 和 2 的元素的计数，然后使用这些值。