二叉树根节点到叶子节点的所有路径和

1,412 阅读3分钟

小知识,大挑战!本文正在参与“程序员必备小知识”创作活动

问题描述

给你一个二叉树的根节点 root ,树中每个节点都存放有一个0 到 9 之间的数字。每条从根节点到叶节点的路径都代表一个数字:例如,从根节点到叶节点的路径 1 -> 2 -> 3 表示数字 123。计算从根节点到叶节点生成的所有数字之和。

示例:

输入:root = [1,2,3]

输出:25

解释:从根到叶子节点路径 1->2 代表数字 12,从根到叶子节点路径 1->3 代表数字13,因此,数字总和 = 12 + 13 = 25。

image.png

分析问题

根据题意,我们可以知道每个节点对应一个数字,等于其父节点对应的数字乘以10再加上该节点的值。只要计算出所有叶子节点对应的数字,然后对其求和,即可得到结果。遍历求出所有叶子节点对应的数字,我们可以采用深度优先搜索和广度优先搜索来实现。

深度优先搜索

从根节点开始遍历每个节点,如果遇到叶子节点,则将叶子节点对应的数字加到结果中。如果遇到的是非叶子节点,则计算该节点的值,然后对其子节点进行递归遍历。

class Solution:
    def sumNumbers(self, root):
        def dfs(root, pre):
            #如果为空,返回0
            if not root:
                return 0
            total = pre * 10 + root.val
            #如果遇到叶子节点,加到结果中,否则递归求解其子节点
            if not root.left and not root.right:
                return total
            else:
                return dfs(root.left, total) + dfs(root.right, total)
        #从根节点开始递归    
        return dfs(root, 0)

该算法的时间复杂度和空间复杂度都是O(n)。

广度优先搜索

使用广度优先搜索来求解此题时,我们需要维护两个队列来分别存储节点和节点对应的数字。开始时我们将根节点和根节点对应的值分别加入两个队列中。如果队列不为空,我们就从两个队列分别取出一个节点和一个数字,进行如下操作:

  • 如果当前节点是叶子节点,我们就将该节点对应的数字加入到结果中。
  • 如果当前节点不是叶子节点,则将当前节点的非空子节点入队,同时根据当前节点对应的数字和其子节点对应的数字计算出子节点代表的数字,并入队。

搜索结束后,就可以得到所有叶节点对应的数字之和。

image.png

image.png

下面我们来看一下代码实现。

import collections
class Solution:
    def sumNumbers(self, root):
        #树为空,直接返回0
        if not root:
            return 0

        total = 0
        #根节点和对应的值都入队
        nodeQueue = collections.deque([root])
        numQueue = collections.deque([root.val])
        #如果队列不为空
        while nodeQueue:
            #出队
            node = nodeQueue.popleft()
            num = numQueue.popleft()
            left, right = node.left, node.right
            #如果是叶子节点,加入结果中
            if not left and not right:
                total += num
            else:
                #否则,计算其叶子节点和其对应的数字,并同时入队
                if left:
                    nodeQueue.append(left)
                    numQueue.append(num * 10 + left.val)
                if right:
                    nodeQueue.append(right)
                    numQueue.append(num * 10 + right.val)

        return total

该算法的时间复杂度和空间复杂度都是O(n)。