二叉树中的指数路径计数

简介

在二叉树中，如果存在一条路径，使得该路径上的所有节点的值之和等于$2^n$，那么我们称这条路径是一个“指数路径”。本题要求给定二叉树，计算其中存在多少条指数路径。

思路

我们知道，给定一条路径，可以通过求出它上面所有节点的值之和，来判断该路径是否为指数路径。但是，常规的二叉树遍历方式并不适合该问题。因此，我们需要借助一些其他的技巧。

由于我们要判断的是指数路径，即路径上节点的值之和等于$2^n$，我们可以将每个节点的值都转化为它们对应的二进制数的指数。例如，节点 $9$ 的值可以转化为 $2^3$。

接下来，我们可以利用前缀和的思想来快速判断是否存在指数路径。我们可以用一个变量 $cur$ 来表示当前遍历到的节点的值之和对应的二进制指数。同时，我们可以用一个哈希表 $mp$ 来记录从根节点到当前遍历的节点的路径上，每个指数在路径上出现的次数。

以节点 $x$ 为起点，节点 $y$ 为终点的路径是否是指数路径，我们可以通过计算 $cur$ 的变化量来判断。假设 $x$ 到 $y$ 的路径上有一段中间节点 $z$，那么从 $x$ 到 $z$ 的路径对应的指数之和为 $cur-k$，从 $x$ 到 $y$ 的路径对应的指数之和为 $cur$。如果这两个指数相差为 $k$，并且从 $x$ 到 $z$ 的路径和对应的指数已经出现在哈希表 $mp$ 中，那么从 $x$ 到 $y$ 的路径就是一条指数路径。

最后，为避免重复计数，我们需要在访问节点 $y$ 后，将 $mp$ 中的所有指数出现次数减一。

代码

class Solution:
    def __init__(self):
        self.count = 0
        self.cur = 0
        self.mp = defaultdict(int)

    def dfs(self, root):
        if not root:
            return
        # 更新哈希表
        self.cur += 1 << root.val
        self.count += self.mp[self.cur - 2**root.val]
        self.mp[self.cur] += 1
        # 递归遍历左右子树
        self.dfs(root.left)
        self.dfs(root.right)
        # 回溯时更新哈希表
        self.mp[self.cur] -= 1
        self.cur -= 1 << root.val

    def countPairs(self, root: Optional[TreeNode], target: int) -> int:
        self.dfs(root)
        return self.count

代码中，我们用了一个 defaultdict 来作为哈希表。默认情况下，哈希表中不存在的键，对应的值会被自动初始化为 0。因此，我们可以避免在更新哈希表时，需要判断指定键是否已存在的情况。