思路详解

问题分析

我们需要计算树中从根节点出发能够到达的所有节点的路径总数. 这里的路径指的是从根节点向下延伸的简单路径 (不包含回头路).

错误思路分析

最初的朴素想法是从底层向上逐层计算每个节点对上层节点的贡献. 具体来说：

1. 每个节点自身贡献1条路径 (只包含自己)
2. 每个节点会向其所有非父节点的上层节点贡献其路径数 但这种做法需要对每个节点向上遍历所有祖先节点, 时间复杂度为$O (n²)$ , 对于$n=3×10⁵$的数据规模显然不可行

正确思路优化

采用动态规划的思想, 结合树的分层处理：

1. 分层处理：首先将树按层次划分, 根节点在第0层, 其子节点在第1层, 以此类推.

2. 路径计数规则：

   • 每个节点至少有1条路径 (只包含自己)
   • 每个节点会向其上层节点贡献其路径数 (因为上层节点可以通过子节点到达这些路径)
   • 但父节点不能通过子节点到达子节点本身 (即不能直接使用子节点的路径数)

3. 关键优化：

   • 从底层向上处理, 维护一个”上一层总和”变量
   • 对于当前层的每个节点：
     • 首先加上上一层所有节点的路径总和 (因为这些节点都可以被当前节点的父节点到达)
     • 然后从父节点的计数中减去当前节点的计数 (消除父节点不能到达当前节点的部分)

4. 状态转移方程：

   • 设$dp[u]$表示以$u$为根的子树中从根到$u$的路径数
   • $dp[u] += 1 + Σdp[v] - Σdp[w]$ ($v$是$u$的下层节点, $w$是$u$的子节点)

复杂度分析

• 时间复杂度：$O (n)$
  • 分层处理需要 $O (n)$ 时间
  • 自底向上处理每层节点也是 $O (n)$
• 空间复杂度：$O (n)$
  • 需要存储树的结构和每层节点信息

完整代码

for _ in range(int(input())):
    # 数据输入
    n = int(input())
    a = map(int, input().split())

    # 下标1对应根节点. 下标0对应0节点, 但是并无0节点, 于是占位.
    # 三个数值分别对应 [父节点, 层数, 路径计数].
    # 显然的, 每一个节点至少有一条路径, 即为本身.
    info = [[0, 0, 0], [0, 0, 1]] 
    level = [[1]] # 第0层有且仅有根节点.
    for i, parent in enumerate(a):
        lev = info[parent][1] + 1 # 当前层数为父节点层数加1.
        info.append([parent, lev, 1]) # 更新节点信息.
        if len(level) <= lev: # 如果当前层数不足则添加一层. (显而易见, 层数只能一次增加1)
            level.append([]) 
        level[lev].append(i + 2) # 在目标层加入点, i从0开始, 但是点id从1开始, 并且不包括根节点, 所以点id应该为i+2.

    last_lsum = 0 # 上一层的总和
    for lst in level[::-1]: # 从最后一层开始遍历.
        level_sum = 0 # 当前层总和
        for i in lst:
            info[i][2] += last_lsum # 将当前点添加上一层的总和.
            level_sum += info[i][2] # 将当前点的计数添加到当前层总和
            if info[i][0] != 1: # 特判父节点为根节点.
                info[info[i][0]][2] -= info[i][2] # 将父节点减少当前节点计数(因为父节点无法到达子节点).

        last_lsum = level_sum # 更新上一层的总和.

    print(info[1][2] % 998244353) # 输出.

# 提交记录: https://codeforces.com/contest/2070/submission/312559631
# 语言: pypy3.10
# Code by @hsn8086