LCA+前缀和 WA#11#15求调

#include <bits/stdc++.h>
using namespace std;
using i64 = long long;
constexpr i64 N = 3e5 + 5, MAXLOG = 18, K = 55, mod = 998244353;
// Graph representation
vector<int> G[N];
i64 dep[N], mul[K], sa[N][K], anc[N][MAXLOG];

i64 ksm(i64 a, i64 b) {
    // Fast exponentiation modulo mod
    i64 res = 1;
    while (b) {
        if (b & 1) res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res;
}

void dfs(int x, int fa) {
    anc[x][0] = fa;
    for (int i = 1; i < MAXLOG; ++i) {
        anc[x][i] = anc[anc[x][i - 1]][i - 1];
    }
    for (int y : G[x]) {
        if (y != fa) {
            dep[y] = dep[x] + 1;
            for (int i = 1; i <= 50; i++) {
                mul[i] = mul[i - 1] * dep[y] % mod;
            }
            for (int i = 1; i <= 50; i++) {
                sa[y][i] = (sa[x][i] + mul[i]) % mod;
            }
            dfs(y, x);
        }
    }
}

int LCA(int x, int y) {
    // Find the Lowest Common Ancestor (LCA) of nodes x and y
    if (dep[x] < dep[y]) swap(x, y);
    for (int i = MAXLOG - 1; i >= 0; --i) {
        if (dep[anc[x][i]] >= dep[y]) {
            x = anc[x][i];
        }
    }
    if (x == y) return x;
    for (int i = MAXLOG - 1; i >= 0; --i) {
        if (anc[x][i] != anc[y][i]) {
            x = anc[x][i];
            y = anc[y][i];
        }
    }
    return anc[x][0];
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    
    // Input number of nodes
    int n;
    cin >> n;
    // Initialize the graph
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    // Start DFS from node 1 (assuming 1 is the root)
    mul[0] = 1; // Initialize the multiplier for depth 0
    dfs(1, 0);
    // Input number of queries
    int m;
    cin >> m;
    // Process each query
    while (m--) {
        int u, v, k;
        cin >> u >> v >> k;
        // Find the LCA of u and v
        int lca = LCA(u, v);
        // Calculate the result using the precomputed powers
        i64 res = (sa[u][k] + sa[v][k] - sa[anc[lca][0]][k] - sa[lca][k] + 2 * mod) % mod;
        cout << res << '\n';
    }
    return 0;
}