二叉树深度遍历终结版

Posted 2021-04-09 算法的秘密

❝

前面已经有3篇文章讲二叉树深度优先遍历，分别是   。今天这道题是之前问题的进阶版。学会它，基本具备与大厂面试官谈笑风生的能力，想想都很兴奋呢！

❞

1 问题描述

给定一个二叉树和一个数字S，找出所有的路径，使得路径上所有元素的和等于S。这里的路径的起点和终点可以是任何节点，但是要遵循从父节点到子节点的方向。

「例子1」给定S=12，和如下二叉树：

输出是：3

因为可以发现：7-5，1-9-2，9-3三个路径都满足要求。

「例子2」给定S=11和如下二叉树：

输出是：2

因为可以发现：7-4和1-10两个路径都满足要求。

2 问题分析

「复杂的问题有时可以从简单的角度进行分析。」

对这道题，咱们只要能找到所有可能的路径，挨个检查就可以了。

那么如何找到所有可能的路径呢？

【下面这段话请仔细品味】

我们可以将所有可能的路径「按照路径终点所在节点」进行分类。结合深度优先遍历，在遍历到每一个节点时，对以该节点为终点的所有路径进行检查就可以了。

当然，为了方便遍历所有可能路径，我们需要记录下访问该节点时的当前路径信息。

3 代码实现

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


def count_paths(root, S):
  return count_paths_recursive(root, S, [])


def count_paths_recursive(currentNode, S, currentPath):
  if currentNode is None:
    return  0

  # 将当前节点加入到当前路径中
  currentPath.append(currentNode.val)
  pathCount, pathSum = 0, 0
  # 检查所有以当前节点为终点的路径
  for i in range(len(currentPath)-1, -1, -1):
    pathSum += currentPath[i]
    # 找到一条符合条件的路径时，计数加1
    if pathSum == S:
      pathCount += 1

  # 递归调用左子树
  pathCount += count_paths_recursive(currentNode.left, S, currentPath)
  # 递归调用右子树
  pathCount += count_paths_recursive(currentNode.right, S, currentPath)

  # 从递归调用栈退出时，将当前节点从当前路径中删除，以便进行回溯
  del currentPath[-1]

  return pathCount

往期精彩回顾