先dfs1求出每个节点的深度,然后预处理出每个节点k次方的深度,50个线段树每个线段树存k次方的值。 求大佬帮忙看看怎么才能不mle

#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
#define int long long
using namespace std;
int n, m;
#define lc u<<1
#define rc u<<1|1
const int N = 3e5 + 7;;
const int M = 998244353;
vector<int>edge[N];
int tim;
int dep[51][N], siz[N], son[N], fa[N], top[N], id[N],nw[53][N];
struct Node {
    int  sum;
}t[53][N << 2];
void dfs1(int u, int father) {
    fa[u] = father; dep[1][u] = dep[1][father] + 1; siz[u] = 1;
    for (auto v : edge[u]) {
        if (v == father)continue;
        dfs1(v, u);
        siz[u] += siz[v];
        if (siz[son[u]] < siz[v])son[u] = v;
    }
}
void dfs2(int u, int tp) {
    top[u] = tp; id[u] = ++tim;
    for (int i = 1; i <= 50; i++) {
        nw[i][tim] = dep[i][u];
    }
        if (!son[u])return;
        dfs2(son[u], tp);
        for (auto v : edge[u]) {
            if (v == fa[u] || v == son[u])continue;
            dfs2(v, v);
        }
}
void pushup(int u,int i) {
    t[i][u].sum = (t[i][lc].sum + t[i][rc].sum)%M;
}
int query_sum(int u, int L, int R,int l,int r,int i) {
    if (L <= l && r <= R) {
        return t[i][u].sum;
    }
    int res = 0;
    int mid = (l + r) >> 1;
    if (L <= mid)
        res = (res + query_sum(lc, L, R,l,mid, i)) % M;
    if (R > mid)
        res = (res + query_sum(rc, L, R,mid+1,r, i)) % M;
    return res;

}
void build(int u, int l, int r,int i) {
    t[i][u] = { nw[i][l] };
    if (l == r)return;
    int m = l + r >> 1;
    build(lc, l, m, i);
    build(rc, m+1, r, i);
    pushup(u, i);
}
int query_tree_sum(int u,int v,int k) {
    int res = 0;
    while (top[u] != top[v]) {
        if (dep[1][top[u]] < dep[1][top[v]]) {
            swap(u, v);
        }
        res = (res+query_sum(1, id[top[u]], id[u],1,n,k))%M;
        u = fa[top[u]];
        }
    if (dep[1][u] < dep[1][v]) swap(u, v);
    res = (res+query_sum(1, id[v], id[u],1,n,k))%M;
    return res;
}
signed main() {
#ifdef LOCAL
    freopen("in.in", "r", stdin);
    freopen("out.out", "w", stdout);
#endif
    //===============
    //===============
end:
    //1 2 2 3 3
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n;
    for (int i = 1; i < n; i++) {
        int x, y; cin >> x >> y;
        edge[x].push_back(y);
        edge[y].push_back(x);
    }
    dfs1(1, 0);
    
    for (int i = 1; i <= n; i++) {
        dep[1][i]--;
    }
    //k i
    for (int i = 2; i <= 50; i++) {
        for (int j = 1; j <= n; j++) {
            dep[i][j] = dep[i - 1][j] * dep[1][j] % M;
        }
    }
    dfs2(1, 1);
    for (int i = 1; i <= 50; i++) {
        build(1, 1, n,i);
    }
    cin >> m;
    for (int i = 1; i <= m; i++) {
        int ii, jj, kk; cin >> ii >> jj >> kk;
        cout << query_tree_sum(ii, jj, kk) << endl;
    }
    return 0;
}