[BJOI2018]求和

题目链接:

点我

题目分析:

观察到\(k\)只有\(50\),时间复杂度完全能承受直接预处理
在每个点上预处理一下从根到这个点的所有节点深度的\(k\)次方和,然后树上差分一下即可

代码:

#include<bits/stdc++.h>
#define N (300000 + 10)
#define ll long long
#define int ll
using namespace std;
inline int read() {
	int cnt = 0, f = 1; char c = getchar();
	while (!isdigit(c)) {if (c == '-') f = -f; c = getchar();}
	while (isdigit(c)) {cnt = (cnt << 3) + (cnt << 1) + (c ^ 48); c = getchar();}
	return cnt * f;
}
const int mod = 998244353;
int n, m, mul[N][35], dep[N], a, b, k;
ll sum[N][55];
int first[N], nxt[N << 1], to[N << 1], tot;
void add(int x, int y) {nxt[++tot] = first[x], first[x] = tot, to[tot] = y;}
ll qpow(ll a, ll b) {
	ll ans = 1;
	for (; b; b >>= 1) {if (b & 1) ans *= a, ans %= mod; a *= a, a %= mod;}
	return ans % mod;
}
void dfs_(int u, int f) {
	dep[u] = dep[f] + 1;
	for (register int i = 1; i <= 50; ++i) sum[u][i] += qpow(dep[u], i);
	for (register int i = first[u]; i; i = nxt[i]) {
		int v = to[i];
		if (v == f) continue;
		mul[v][0] = u;
		for (register int j = 1; j <= 50; ++j) sum[v][j] = (sum[v][j] + sum[u][j]) % mod;
		dfs_(v, u);
	}
}
void pre_work() {
	for (register int j = 1; j <= 23; ++j) 
		for (register int i = 1; i <= n; ++i) mul[i][j] = mul[mul[i][j - 1]][j - 1];
}
int lca(int x, int y) {
	if (dep[x] < dep[y]) swap(x, y);
	for (register int i = 23; i >= 0; --i) 
		if (dep[mul[x][i]] >= dep[y]) x = mul[x][i];
	if (x == y) return x;
	for (register int i = 23; i >= 0; --i) 
		if (mul[x][i] != mul[y][i]) x = mul[x][i], y = mul[y][i];
	return mul[x][0];
}
signed main() {
	n = read();
	for (register int i = 1; i < n; ++i) {
		a = read(), b = read();
		add(a, b), add(b, a);
	}
	dep[0] = -1;
	dfs_(1, 0); pre_work();
	m = read();
	for (register int i = 1; i <= m; ++i) {
		a = read(), b = read(), k = read();
		printf("%lld\n", ((sum[a][k] + sum[b][k]) % mod - sum[lca(a, b)][k] % mod - sum[mul[lca(a, b)][0]][k] % mod + mod * 2) % mod);
	}
	return 0;
}
posted @ 2019-11-12 08:19  kma_093  阅读(98)  评论(0编辑  收藏  举报