树结构算法：深度k次方和与LCA应用

1. 问题背景与核心思路

这道题目来自洛谷P4427 [BJOI2018]求和，属于树结构相关算法题。题目要求我们处理一棵以1号节点为根的树，支持多次查询：给定两个节点u、v和一个整数k，求u到v路径上所有节点深度的k次方之和。这里的深度指的是节点到根节点的距离（边数）。

我第一次看到这个问题时，最直接的暴力解法就是对每次查询都从u走到v，累加路径上每个节点的depth^k。但这样时间复杂度是O(Q*N)，对于N=3e5、Q=3e5的数据规模显然会超时。于是我们需要更聪明的预处理方法。

核心突破点在于发现depth^k的前缀和性质。想象一下在一维数组中，我们计算区间和可以用前缀和数组相减。在树上是否也存在类似的"前缀和"？答案是肯定的——我们可以预处理每个节点到根节点路径上所有depth^k的和，然后利用最近公共祖先(LCA)来拆分路径。

2. 算法设计与数学原理

2.1 树上前缀和定义

我们定义s[u][k]表示从根节点到u节点路径上所有节点depth的k次方之和。例如对于k=2：

• 根节点1：s[1][2] = 0² = 0
• 节点2（depth=1）：s[2][2] = 0² + 1² = 1
• 节点3（depth=2）：s[3][2] = 0² + 1² + 2² = 5

这个预处理可以通过一次DFS完成，时间复杂度O(N*K)，其中K是最大的k值（题目中k≤50）。

2.2 路径求和公式推导

对于查询u到v路径上的depth^k和，我们可以将其拆分为四部分：

1. u到根的路径和：s[u][k]
2. v到根的路径和：s[v][k]
3. LCA到根的路径和：s[LCA][k]
4. LCA父节点到根的路径和：s[fa[LCA][0]][k]

最终的求和公式为：
ans = s[u][k] + s[v][k] - s[LCA][k] - s[fa[LCA][0]][k]

这个公式的原理是：u到v的路径可以看作u→LCA→v，而我们加上了u和v到根的路径，减去了LCA及其父节点到根的路径（因为它们被重复计算了）。

2.3 倍增法求LCA

为了高效查询任意两节点的LCA，我们使用倍增法预处理：

• fa[u][i]表示u节点的2^i级祖先
• 预处理时采用DFS，递推公式：fa[u][i] = fa[fa[u][i-1]][i-1]
• 查询时先将两个节点调整到同一深度，然后一起向上跳

这样预处理O(NlogN)，每次查询O(logN)。

3. 代码实现细节

3.1 数据结构定义

cpp复制const int MAXN = 3e5+10, MAXK = 55, MOD = 998244353;
vector<int> tree[MAXN];  // 树的邻接表表示
int depth[MAXN];         // 节点深度
long long s[MAXN][MAXK]; // 前缀和数组
int fa[MAXN][20];        // 倍增数组

3.2 预处理阶段

cpp复制void dfs(int u, int parent) {
    depth[u] = depth[parent] + 1;
    fa[u][0] = parent;
    
    // 计算所有k次方的前缀和
    long long pow = 1;
    for(int k=1; k<=50; k++) {
        pow = pow * depth[u] % MOD;
        s[u][k] = (s[parent][k] + pow) % MOD;
    }
    
    // 倍增预处理
    for(int i=1; i<20; i++)
        fa[u][i] = fa[fa[u][i-1]][i-1];
        
    for(int v : tree[u]) {
        if(v != parent) dfs(v, u);
    }
}

3.3 LCA查询实现

cpp复制int lca(int u, int v) {
    if(depth[u] < depth[v]) swap(u, v);
    
    // 将u提到与v同一深度
    for(int i=19; i>=0; i--)
        if(depth[fa[u][i]] >= depth[v])
            u = fa[u][i];
    
    if(u == v) return u;
    
    // 一起向上跳
    for(int i=19; i>=0; i--)
        if(fa[u][i] != fa[v][i])
            u = fa[u][i], v = fa[v][i];
    
    return fa[u][0];
}

3.4 查询处理

cpp复制long long query(int u, int v, int k) {
    int ancestor = lca(u, v);
    return (s[u][k] + s[v][k] - s[ancestor][k] - s[fa[ancestor][0]][k] + 2*MOD) % MOD;
}

注意：最后要加上2*MOD再取模，因为两个减法可能导致结果为负。

4. 复杂度分析与优化

4.1 时间复杂度

• 预处理阶段：

  • DFS遍历：O(N)
  • 前缀和计算：O(N*K) (K=50)
  • 倍增预处理：O(NlogN)
  • 总计：O(N50 + NlogN) ≈ O(3e550) = 1.5e7

• 查询阶段：

  • 每次LCA查询：O(logN) ≈ 20
  • 前缀和查询：O(1)
  • 总计：O(Q*20) ≈ 6e6

完全在合理范围内。

4.2 空间复杂度

• 前缀和数组：3e5508 ≈ 120MB
• 倍增数组：3e5204 ≈ 24MB
• 其他：可忽略
• 总计：约150MB，远小于通常的512MB限制

5. 常见问题与调试技巧

5.1 负数取模问题

在查询函数中，我们可能会遇到：
s[u][k] + s[v][k] - s[LCA][k] - s[fa[LCA][0]][k] 为负数的情况。解决方法：

1. 加上足够大的MOD倍数（通常是2*MOD）
2. 再取模保证结果非负

5.2 边界条件处理

• 根节点1的父节点应设为0，且depth[0]=-1
• 这样depth[1]=0，符合题目定义
• 在LCA查询时要注意u或v已经是祖先的情况

5.3 性能优化

• 使用快速读入（对于3e5量级的输入很重要）
• 将k的循环上限设为输入的最大k而非固定的50
• 使用邻接表而非邻接矩阵存储树结构

6. 实际编码中的经验分享

在实现这个算法时，我踩过几个坑值得分享：

1. 深度定义一致性：题目中说"根的深度为0"，但有些人习惯设为1。这会影响所有depth^k的计算，必须统一。我选择严格遵循题意，将根深度设为0。

2. 幂次计算优化：最初我每个k都重新计算depth[u]^k，这样会有大量重复计算。优化方法是：

cpp复制long long pow = 1;
for(int k=1; k<=50; k++) {
    pow = pow * depth[u] % MOD;
    s[u][k] = (s[parent][k] + pow) % MOD;
}

这样每个节点只需O(K)时间而非O(KlogK)。

3. LCA的二进制跳跃：在实现倍增法时，跳跃要从大到小尝试：

cpp复制for(int i=19; i>=0; i--)  // 不是从0到19！
    if(depth[fa[u][i]] >= depth[v])
        u = fa[u][i];

这个顺序很重要，能保证我们跳最远的合法距离。

4. 内存占用：s[MAXN][MAXK]数组看似很大，但实际计算：

• 3e5节点 × 50k × 8字节 ≈ 120MB
  这在编程竞赛中是可接受的，但要注意不要爆空间。

这道题综合考察了树的前缀和、LCA、模运算等多个知识点，是一道质量很高的树结构练习题。通过这道题，我对树上路径问题的处理有了更深的理解——很多线性结构中的技巧（如前缀和）经过适当变形，也能应用于树结构。