[BJOI2018]求和

嘟嘟嘟

首先看到k只有50，那么就可以开一个数组预处理出来。

sum[u][k]表示节点u到根节点所有节点深度的k次方和，dfs一遍就都搞出来了，预处理复杂度O(n * 50)（快速幂复杂度不计了）。

查询就是lca复杂度，对于路径(x, y)，令z = lca(x, y)，则ans(x, y)  = sum[x][k] + sum[y][k] - sum[z][k] - sum[fa[z]][k]。注意取模别出负数。

#include<cstdio>
#include<iostream>
#include<cmath>
#include<algorithm>
#include<cstring>
#include<cstdlib>
#include<cctype>
#include<vector>
#include<stack>
#include<queue>
using namespace std;
#define enter puts("") 
#define space putchar(' ')
#define Mem(a, x) memset(a, x, sizeof(a))
#define rg register
typedef long long ll;
typedef double db;
const int INF = 0x3f3f3f3f;
const db eps = 1e-8;
const int maxn = 3e5 + 5;
const ll mod = 998244353;
inline ll read()
{
      ll ans = 0;
    char ch = getchar(), last = ' ';
    while(!isdigit(ch)) {last = ch; ch = getchar();}
    while(isdigit(ch)) {ans = ans * 10 + ch - '0'; ch = getchar();}
    if(last == '-') ans = -ans;
    return ans;
}
inline void write(ll x)
{
    if(x < 0) x = -x, putchar('-');
    if(x >= 10) write(x / 10);
    putchar(x % 10 + '0');
}

int n, m;
struct Edge
{
    int nxt, to;
}e[maxn << 1];
int head[maxn], ecnt = -1;
void addEdge(int x, int y)
{
    e[++ecnt] = (Edge){head[x], y};
    head[x] = ecnt;
}

ll quickpow(ll a, int b)
{
    a %= mod;
    ll ret = 1;
    for(; b; a = a * a % mod, b >>= 1) 
        if(b & 1) ret = ret * a % mod;
    return ret;
}

int dis[maxn];
ll sum[maxn][55];
void dfs(int now, int f)
{
    for(int i = head[now]; i != -1; i = e[i].nxt)
    {
        if(e[i].to == f) continue;
        dis[e[i].to] = dis[now] + 1;
        for(int j = 1; j <= 50; ++j) sum[e[i].to][j] = (sum[now][j] + quickpow(dis[e[i].to], j)) % mod;
        dfs(e[i].to, now);
    }
}

const int N = 22;
int fa[maxn][25];
void dfs2(int now, int f)
{
    for(int i = 1; i <= N; ++i)
        fa[now][i] = fa[fa[now][i - 1]][i - 1];
    for(int i = head[now]; i != -1; i = e[i].nxt)
    {
        if(e[i].to == f) continue;
        fa[e[i].to][0] = now;
        dfs2(e[i].to, now);
    }
}
int lca(int x, int y)
{
    if(dis[x] < dis[y]) swap(x, y);
    for(int i = N; i >= 0; --i)
        if(dis[x] - (1 << i) >= dis[y]) x = fa[x][i];
    if(x == y) return x;
    for(int i = N; i >= 0; --i)
        if(fa[x][i] != fa[y][i]) x = fa[x][i], y = fa[y][i];
    return fa[x][0];
}

ll solve(int x, int y, int k)
{
    int z = lca(x, y);
    return (sum[x][k] + sum[y][k] - (sum[z][k] + sum[fa[z][0]][k]) % mod + mod) % mod;
}

int main()
{
    Mem(head, -1);
    n = read();
    for(int i = 1; i < n; ++i)
    {
        int x = read(), y = read();
        addEdge(x, y); addEdge(y, x);
    }
      dfs(1, 0);
      dfs2(1, 0);
      m = read();
      for(int i = 1; i <= m; ++i)
      {
          int x = read(), y = read(), k = read();
          write(solve(x, y, k)), enter;
      }
      return 0;
}