题解：洛谷 P4427 [BJOI2018] 求和

【题目来源】

洛谷：[P4427 BJOI2018] 求和 - 洛谷

【题目描述】

master 对树上的求和非常感兴趣。他生成了一棵有根树，并且希望多次询问这棵树上一段路径上所有节点深度的 \(k\) 次方和，而且每次的 \(k\) 可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给了 pupil，但 pupil 并不会这么复杂的操作，你能帮他解决吗？

【输入】

第一行包含一个正整数 \(n\)，表示树的节点数。

之后 \(n-1\) 行每行两个空格隔开的正整数 \(i, j\)，表示树上的一条连接点 \(i\) 和点 \(j\) 的边。

之后一行一个正整数 \(m\)，表示询问的数量。

之后每行三个空格隔开的正整数 \(i, j, k\)，表示询问从点 \(i\) 到点 \(j\) 的路径上所有节点深度的 \(k\) 次方和。由于这个结果可能非常大，输出其对 \(998244353\) 取模的结果。

树的节点从 \(1\) 开始标号，其中 \(1\) 号节点为树的根。

【输出】

对于每组数据输出一行一个正整数表示取模后的结果。

【输入样例】

5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45

【输出样例】

33
503245989

【算法标签】

《洛谷 P4427 求和》 #倍增# #深度优先搜索,DFS# #最近公共祖先,LCA# #树链剖分# #前缀和# #各省省选# #北京# #O2优化# #2018#

【代码详解】

#include <bits/stdc++.h>
using namespace std;

#define int long long  // 使用长整型
const int N = 300005, M = 2 * N;  // N: 最大节点数, M: 最大边数
const int mod = 998244353;        // 模数
int n, m;                         // n: 节点数, m: 查询次数
int tot, to[M], ne[M], h[N];      // 链式前向星存储树结构
int fa[N][22];                    // 倍增数组，fa[u][i]表示u的2^i级祖先
int dep[N];                       // 节点深度数组
int mi[60];                       // 临时存储深度幂次
int s[N][60];                     // s[v][j]表示根到v路径上节点深度的j次幂和

// 添加边到链式前向星
void add(int a, int b)
{
    to[++tot] = b;                // 存储终点
    ne[tot] = h[a];               // 存储下一条边
    h[a] = tot;                   // 更新头指针
}

// 深度优先搜索预处理节点信息
void dfs(int u, int f)
{
    // 预处理倍增数组
    for (int i = 1; i <= 20; i++)
        fa[u][i] = fa[fa[u][i - 1]][i - 1];  // 递推计算2^i级祖先
    
    // 遍历所有邻接节点
    for (int i = h[u]; i; i = ne[i])
    {
        int v = to[i];
        if (v == f) continue;     // 跳过父节点
        
        // 初始化子节点信息
        fa[v][0] = u;             // 直接父节点
        dep[v] = dep[u] + 1;      // 深度+1
        
        // 计算深度幂次
        mi[0] = 1;
        for (int j = 1; j <= 50; j++)
            mi[j] = mi[j - 1] * dep[v] % mod;
        
        // 计算幂次前缀和
        for (int j = 1; j <= 50; j++)
            s[v][j] = (mi[j] + s[u][j]) % mod;
        
        // 递归处理子节点
        dfs(v, u);
    }
}

// 倍增法求最近公共祖先
int lca(int u, int v)
{
    // 保证u是较深的节点
    if (dep[u] < dep[v]) swap(u, v);
    
    // 将u提升到与v相同深度
    for (int i = 20; i >= 0; i--)
        if (dep[fa[u][i]] >= dep[v])
            u = fa[u][i];
    
    if (u == v) return v;         // 如果已经是同一个节点
    
    // 同时向上跳跃
    for (int i = 20; i >= 0; i--)
        if (fa[u][i] != fa[v][i])
            u = fa[u][i], v = fa[v][i];
    
    return fa[u][0];              // 返回LCA
}

signed main()
{
    // 输入树结构
    cin >> n;
    for (int i = 1; i < n; i++)
    {
        int a, b;
        cin >> a >> b;
        add(a, b);                // 无向图添加双向边
        add(b, a);
    }
    
    // 预处理树信息
    mi[0] = 1;
    dfs(1, 0);                    // 从根节点1开始DFS
    
    // 处理查询
    cin >> m;
    for (int i = 1; i <= m; i++)
    {
        int u, v, k;
        cin >> u >> v >> k;
        int l = lca(u, v);        // 求LCA
        
        // 计算路径上的深度k次幂和
        int ans = (s[u][k] + s[v][k] - s[l][k] - s[fa[l][0]][k] + 2 * mod) % mod;
        cout << ans << endl;
    }
    
    return 0;
}

【运行结果】

5
1 2
1 3
2 4
2 5
2
1 4 5
33
5 4 45
503245989