题解:洛谷 P4427 求和
【题目来源】
【题目描述】
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
【核心思想】
-
问题分析:给定一棵 \(n\) 个节点的有根树(根为 \(1\)),节点深度定义为到根路径上的边数。有 \(m\) 次询问,每次询问给出两个节点 \(u, v\) 和一个整数 \(k\)(\(1 \le k \le 50\)),要求计算从 \(u\) 到 \(v\) 的路径上所有节点深度的 \(k\) 次方之和,并对 \(998244353\) 取模。由于 \(n, m\) 可达 \(3\times 10^5\),\(k\) 的上限为 \(50\),不能对每次询问进行路径遍历,需利用树上前缀和将路径和转化为端点前缀和的组合。
-
算法选择:
- 树上前缀和(深度幂次和):预处理每个节点 \(u\) 到根节点路径上所有节点深度的 \(1\sim 50\) 次方和,记为 \(sum[u][k]\)。则从根到 \(u\) 的深度 \(k\) 次方和满足递推 \(sum[u][k] = sum[fa[u]][k] + dep[u]^k\)。
- 倍增法求 LCA:预处理每个节点的 \(2^i\) 级祖先和深度,用于 \(O(\log n)\) 求最近公共祖先。
- 路径和公式:对于路径 \(u \to v\),设其 LCA 为 \(l\),则路径上深度 \(k\) 次方和为 \(sum[u][k] + sum[v][k] - sum[l][k] - sum[fa[l][0]][k]\)(模意义下),该公式等价于标准路径点权公式 \(sum[u]+sum[v]-2sum[l]+val[l]\)。
-
关键步骤:
- 读入与建图:读入 \(n\) 及 \(n-1\) 条边,建立无向邻接表,根节点为 \(1\)。
- DFS 预处理:
- 计算每个节点的深度 \(dep[u]\)(根深度为 \(0\))。
- 预处理倍增祖先数组 \(fa[u][i]\)(\(2^i\) 级祖先)。
- 计算 \(dep[u]\) 的 \(1\sim 50\) 次幂(模 \(mod\)),并利用父节点的前缀和递推得到 \(sum[u][k]\)。
- 查询处理:
- 读入 \(u, v, k\),用倍增法求 \(l = LCA(u, v)\)。
- 计算答案 \(ans = (sum[u][k] + sum[v][k] - sum[l][k] - sum[fa[l][0]][k]) \bmod mod\)(注意处理负数,加模数后取模)。
- 输出 \(ans\)。
-
时间/空间复杂度:
- 时间复杂度:预处理 \(O(n \cdot 50 + n \log n)\),每次查询 \(O(\log n)\),总复杂度 \(O((n+m)\log n + 50n)\)。
- 空间复杂度:\(O(n \cdot 50 + n \log n)\),存储前缀和数组和倍增祖先表。
-
树上前缀和 + LCA 的核心思想:
- 前缀和转化:将路径上的点权和转化为端点处前缀和的组合,是树上路径统计问题的通用技巧。本题中每个 \(k\) 值独立,预处理 \(50\) 个前缀和数组,使得查询只需 \(O(1)\) 次前缀和运算。
- LCA 的必要性:路径 \(u \to v\) 的转折点为 LCA,通过 LCA 将路径拆分为 \(u \to l\) 和 \(l \to v\) 两段,再利用前缀和差得到每段的和。
- 幂次可分离性:因为询问的 \(k\) 不同,但 \(k\) 的范围很小(\(\le 50\)),预处理所有可能 \(k\) 的前缀和,使查询与 \(k\) 无关地快速回答。
- 模运算处理:所有加法和减法均在模 \(998244353\) 下进行,确保结果不溢出。
- 适用于树上路径点权统计、多次查询且权值可分离(如多项式、幂次)的场景。
【算法标签】
普及+ #树上前缀和
【代码详解】
#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
浙公网安备 33010602011771号