洛谷P4427 [BJOI2018]求和

\(\Large\textbf{Description: } \large{一颗n个节点的树，m次询问，每次查询点i到点j的路径上所有节点点深度的k次方的和并对998244353取模(1\leq n,m \leq 300000,1\leq k\leq 50)。}\\\)

\(\Large\textbf{Solution: } \large{一开始看到这道题并没有思路，但是注意到k很小，所以我们可以预处理出每个节点到根节点1的路径上点的1到50次方的和，然后每次O(1)查询即可。\\}\)

\(\Large\textbf{Code: }\\\)

#include <cstdio>
#include <algorithm>
#define LL long long
#define gc() getchar()
#define rep(i, a, b) for (int i = (a); i <= (b); ++i)
using namespace std;
const int N = 3e5 + 5;
const int p = 998244353;
int n, m, cnt, head[N], son[N], dep[N], size[N], top[N], fa[N];
LL dis[N][52];

struct Edge {
	int to, next;	
}e[N];	

inline int read() {
	char ch = gc();
	int ans = 0;
	while (ch > '9' || ch < '0') ch = gc();
	while (ch >= '0' && ch <= '9') ans = (ans << 1) + (ans << 3) + ch - '0', ch = gc();
	return ans;	
}

inline void add(int x, int y) {
	e[++cnt].to = y;
	e[cnt].next = head[x];
	head[x] = cnt;
}

inline void dfs1(int x, int y) {
	int Max = 0;
	LL now = 0;
	size[x] = 1;
	fa[x] = y;
	dep[x] = dep[y] + 1;
	now = dep[x];
	rep(i, 1, 50) dis[x][i] = (dis[y][i] + now) % p, now = (now * dep[x]) % p;
	for (int i = head[x]; i ; i = e[i].next) {
		int u = e[i].to;
		dfs1(u, x);
		size[x] += size[u];
		if (size[u] > Max) son[x] = u, Max = size[u];
	}	
}

inline void dfs2(int x, int tp) {
	top[x] = tp;
	if (!son[x]) return;
	dfs2(son[x], tp);
	for (int i = head[x]; i ; i = e[i].next) {
		int u = e[i].to;
		if (u == son[x]) continue;
		dfs2(u, u);
	}	
} 

inline int lca(int x, int y) {
	while (top[x] != top[y]) {
		if (dep[top[x]] < dep[top[y]]) swap(x, y);
		x = fa[top[x]];
	}
	return dep[x] < dep[y] ? x : y;
}

int main() {
	n = read();
	int x, y, k;
	rep(i, 2, n) { x = read(), y = read(); if (x > y) swap(x, y); add(x, y); }
	dep[0] = -1;
	dfs1(1, 0);
	dfs2(1, 1);
	m = read();
	while (m--) {
		x = read(), y = read(), k = read();
		int l = lca(x, y);
		printf("%lld\n", (dis[x][k] + 2 * p - dis[l][k] - dis[fa[l]][k] + dis[y][k]) % p);
	}
	return 0;
}