P4427 [BJOI2018] 求和

解题思路与代码注释

解题思路

这道题目要求计算树中路径上所有节点深度的k次方和。关键在于如何高效处理大量查询，每个查询可能有不同的k值。

核心思路：

1. 预处理深度和幂次和：

   • 使用DFS计算每个节点的深度

   • 预处理每个节点到根节点路径上所有节点的深度k次方和（s[x][k]），其中k的范围是1到50

2. LCA（最近公共祖先）计算：

   • 使用倍增法预处理每个节点的2^i级祖先

   • 通过LCA算法快速找到任意两节点的最近公共祖先

3. 路径求和公式：

   • 对于查询u到v的路径，其和为：s[u][k] + s[v][k] - s[lca][k] - s[father(lca)][k]

   • 这个公式类似于树上前缀和的差分操作

代码注释

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 3e5 + 10, mod = 998244353;
int n, m;
vector<int> g[N];  // 邻接表存储树
int dep[N], f[N][26];  // dep存储深度，f是倍增数组
ll s[N][60];  // s[x][k]表示从根到x路径上所有节点深度的k次方和

// DFS预处理深度、倍增数组和幂次和
void dfs(int x, int fa) {
    f[x][0] = fa;  // x的父节点是fa
    dep[x] = dep[fa] + 1;  // 计算深度
    
    // 预处理倍增数组
    for(int i = 1; i <= 20; i++) {
        int y = f[x][i - 1];
        f[x][i] = f[y][i - 1];
    }
    
    // 预处理幂次和
    ll tmp = 1;  // 计算dep[x]^k
    for(int i = 1; i <= 50; i++) {
        tmp = (tmp * dep[x]) % mod;  // 计算dep[x]^i
        s[x][i] = (s[fa][i] + tmp) % mod;  // 前缀和
    }
    
    // 递归处理子节点
    for(int i = 0; i < g[x].size(); i++) {
        int y = g[x][i];
        if(y != fa) dfs(y, x);
    }
}

// LCA算法（倍增法）
int lca(int x, int y) {
    // 确保x是较深的节点
    if(dep[x] < dep[y]) swap(x, y);
    
    // 将x提升到与y同一深度
    for(int i = 20; i >= 0; i--) {
        int fx = f[x][i];
        if(dep[fx] >= dep[y]) x = fx;
    }
    if(x == y) return x;
    
    // 同时提升x和y
    for(int i = 20; i >= 0; i--) {
        int fx = f[x][i];
        int fy = f[y][i];
        if(fx != fy) {
            x = fx;
            y = fy;
        }
    }
    return f[x][0];
}

int main() {
    cin >> n;
    // 读入树结构
    for(int i = 1; i < n; i++) {
        int x, y; scanf("%d%d", &x, &y);
        g[x].push_back(y);
        g[y].push_back(x);
    }
    
    dep[0] = -1;  // 根节点的父节点深度设为-1
    dfs(1, 0);  // 从根节点开始预处理
    
    cin >> m;
    while(m--) {
        int x, y, k;
        scanf("%d%d%d", &x, &y, &k);
        int rt = lca(x, y);  // 找到最近公共祖先
        
        // 计算路径和公式
        printf("%lld\n", ((s[x][k] + s[y][k]) % mod - (s[rt][k] + s[f[rt][0]][k]) % mod + mod) % mod);
    }
    return 0;
}

路径求和公式解释

公式：ans = (s[x][k] + s[y][k] - s[lca][k] - s[father(lca)][k]) % mod

这个公式的原理是：

1. s[x][k]包含从根到x路径上所有节点的深度k次方和

2. s[y][k]包含从根到y路径上所有节点的深度k次方和

3. 这两者相加会重复计算从根到lca路径上的节点两次

4. 减去s[lca][k]和s[father(lca)][k]是为了：

   • 减去一次重复计算的从根到lca的路径

   • 因为lca节点本身应该只被计算一次

这类似于树上的前缀和差分操作，但不是简单的树上差分，而是利用了前缀和的性质来高效计算路径和。

复杂度分析

1. 预处理：

   • DFS：O(n)

   • 预处理幂次和：O(n*50)

   • 总预处理复杂度：O(n)

2. 查询处理：

   • 每次LCA查询：O(logn)

   • 每次路径和计算：O(1)

   • 总查询复杂度：O(mlogn)