形成 AP 的二叉树中根到叶路径的计数

概述

给定一个二叉树，求其中所有根到叶子节点路径形成等差数列的数量。

即，对于一个路径 $p$，如果其节点值从 $p_1$ 开始，每两个相邻的节点的值之差相等，则称该路径形成等差数列。

示例

比如下面这棵树：

      1
    /   \
   2     3
  / \   / \
 4   5 6   7

其中形成等差数列的路径有：

• $[1, 2, 4]$
• $[1, 2, 5]$
• $[1, 3, 6]$
• $[1, 3, 7]$

故该二叉树中共有 $4$ 条满足条件的路径。

思路

遍历二叉树的同时，对于每一个节点，维护一个 dict，存储以该节点为结尾的所有等差数列。

具体来说，遍历到当前节点 $n$ 时，假设其父节点为 $p$，则遍历 $p$ 时得到的所有等差数列中，有些可以继续往下延伸，有些则不行。要判断一个等差数列是否可以延伸，只需要判断它与 $n$ 的差值是否与之前的相等即可。

遍历结束后，将所有满足条件的等差数列求和即可。

代码

class Solution:
    def countArithmeticSlices(self, root: TreeNode) -> int:
        count = 0
        stack = [(root, {})]
        while stack:
            node, diffs = stack.pop()
            for prev_diff, prev_count in diffs.items():
                curr_diff = node.val - prev_diff
                curr_count = diffs.get(curr_diff, 0)
                count += prev_count
                diffs[curr_diff] = curr_count + prev_count
            diffs[node.val] = diffs.get(node.val, 0) + 1
            if node.left:
                stack.append((node.left, diffs.copy()))
            if node.right:
                stack.append((node.right, diffs.copy()))
        return count

其中，stack 中存储了待遍历的节点及其对应的等差数列信息，diffs 是一个 dict，存储以当前节点为结尾的所有等差数列的差值及其出现次数。stack.pop() 取出一个节点时，会将其对应的 diffs 传递给子节点。由于 diffs 是可变对象，所以需要用 .copy() 方法复制一份。

在遍历子节点前，需要先将当前节点对应的 diffs 更新一下。

最后统计所有满足条件的等差数列即可。