0%

P4427 [BJOI2018]求和

发表于 更新于 分类于

#题意简述

求一棵树上u点到v点路径上所有点权的k次方和。

#简要做法

因为k很小(k<=50)，dfs的时候预处理好每个点的k次方和到根的k次方前缀和，然后打一遍倍增求lca的板子就好了。答案为sum[u][k] + sum[v][k] - sum[Lca][k] - sum[fa[Lca][0]][k]。

此外倍增求lca的同时向上找一定要倒着搜，要不然会跳少。（害人debug半天）

#参考代码

#pragma GCC optimize (3, "Ofast", "inline")
#include <bits/stdc++.h>
#define ll long long
using namespace std;
const ll mod = 998244353;
int n, m;
const int maxn = 300005;
vector <int> G[maxn];
int fa[maxn][20], lim, dep[maxn];
ll sum[maxn][51];
void dfs(int x) {
    for(int i = 0; i < G[x].size(); ++i) {
        int to = G[x][i];
        if(to == fa[x][0]) continue;
        fa[to][0] = x;
        dep[to] = dep[x] + 1;
        ll tmp = dep[to], now = 1;
        for(int j = 1; j <= 50; ++j) {
            now = now * tmp % mod;
            sum[to][j] = (sum[x][j] + now) % mod;
        }
        dfs(to);
    }
}
void init() {
    lim = log2(n) + 1;
    dfs(1);
    for(int i = 0; i < lim; ++i)
        for(int j = 1; j <= n; ++j)
            fa[j][i + 1] = fa[fa[j][i]][i];
}
int lca(int x, int y) {
    if(dep[x] > dep[y])	swap(x , y);
    for(register int i = 0; i <= lim; ++i)
        if((dep[y] - dep[x] >> i) & 1) y = fa[y][i];
    if(x == y) return x;
    for(register int i = lim; i >= 0; --i)
        if(fa[x][i] != fa[y][i]) {
            x = fa[x][i];
            y = fa[y][i];
        }
    return fa[x][0];
}
inline int query(int u, int v, int k) {
    int Lca = lca(u , v);
    int ret = (sum[u][k] + sum[v][k] - sum[Lca][k] - sum[fa[Lca][0]][k]) % mod;
    return (ret + mod) % mod;
}
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr), cout.tie(nullptr);
    cin >> n;
    int u, v;
    for(int i = 1; i < n; ++i) {
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    init();
    cin >> m;
    int k;
    while(m--) {
        cin >> u >> v >> k;
        cout << query(u, v, k) << endl;
    }
    return 0;
}