JZOJ 5691. 【GDOI2018Day2模拟4.25】求和

Description

master对树上的求和非常感兴趣。他生成了一棵有根树，并且希望多次询问这棵树上一段路径上所有节点深度的k次方和，而且每次的k可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给了 Pupil，但Pupil并不会这么复杂的操作，你能帮他解决吗？

Input

从文件 sum.in 中读入数据。
第一行包含一个正整数 n，表示树的节点数。
之后 n−1 行每行两个空格隔开的正整数 i, j，表示树上的一条连接点 i 和点 j 的边。
之后一行一个正整数 m，表示询问的数量。
之后每行三个空格隔开的正整数 i, j,k，表示询问从点 i 到点 j 的路径上所有节点深度的 k 次方和。由于这个结果可能非常大，输出其对 998244353 取模的结果。
树的节点从 1 开始标号，其中 1 号节点为树的根。

Output

输出到文件 sum.out 中。
对于每组数据输出一行一个正整数表示取模后的结果。

Sample Input

5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45

Sample Output

33
503245989

Data Constraint

对于 30% 的数据，1 ≤ n,m ≤ 100。
对于 60% 的数据，1 ≤ n,m ≤ 1000。
对于 100% 的数据，1 ≤ n,m ≤ 300000,1 ≤ k ≤ 50。

【提示】

数据规模较大，请注意使用较快速的输入输出方式。

Hint

以下用 d(i) 表示第 i 个节点的深度。
对于样例中的树，有 d(1) = 0,d(2) = 1,d(3) = 1,d(4) = 2,d(5) = 2。
因此第一个询问答案为 (2^5 + 1^5 + 0^5 ) mod 998244353 = 33，第二个询问答案为(2^45 + 1^45 + 2^45 ) mod 998244353 = 503245989。

Solution

• 好不容易的签到题……

• 直接倍增+LCA，预处理深度k次方的前缀和。

• 询问时求出LCA，因为深度 $d_i$ 是连续的，我们直接利用前缀和计算即可。

• 时间复杂度 $O(N\ log\ N)$ 。

Code

#include<cstdio>
#include<algorithm>
#include<cmath>
#include<cctype>
using namespace std;
typedef long long LL;
const int N=3e5+5,mo=998244353;
int tot;
int first[N],nex[N<<1],en[N<<1];
int dep[N],q[N],f[N][19],g[N][51];
inline int read()
{
    int X=0,w=0; char ch=0;
    while(!isdigit(ch)) w|=ch=='-',ch=getchar();
    while(isdigit(ch)) X=(X<<1)+(X<<3)+(ch^48),ch=getchar();
    return w?-X:X;
}
inline void write(int x)
{
    if(x>9) write(x/10);
    putchar(x%10+'0');
}
inline void insert(int x,int y)
{
    nex[++tot]=first[x];
    first[x]=tot;
    en[tot]=y;
}
void dfs(int x)
{
    dep[x]=dep[f[x][0]]+1;
    for(int i=first[x];i;i=nex[i])
        if(en[i]^f[x][0])
        {
            f[en[i]][0]=x;
            dfs(en[i]);
        }
}
inline int ksm(int x,int y)
{
    int s=1;
    while(y)
    {
        if(y&1) s=(LL)s*x%mo;
        x=(LL)x*x%mo;
        y>>=1;
    }
    return s;
}
inline int getlca(int x,int y)
{
    if(dep[x]<dep[y]) swap(x,y);
    for(int i=log2(dep[x]);i>=0;i--)
        if(dep[f[x][i]]>=dep[y]) x=f[x][i];
    if(x==y) return x;
    for(int i=log2(dep[x]);i>=0;i--)
        if(f[x][i]^f[y][i]) x=f[x][i],y=f[y][i];
    return f[x][0];
}
int main()
{
    freopen("sum.in","r",stdin);
    freopen("sum.out","w",stdout);
    int n=read();
    for(int i=1;i<n;i++)
    {
        int x=read(),y=read();
        insert(x,y);
        insert(y,x);
    }
    int l=0,r=q[1]=dep[1]=1;
    while(l<r)
    {
        int x=q[++l];
        for(int i=first[x];i;i=nex[i])
            if(en[i]^f[x][0])
            {
                f[en[i]][0]=x;
                dep[en[i]]=dep[x]+1;
                q[++r]=en[i];
            }
    }
    //dfs(1);
    for(int j=1;j<=50;j++)
        for(int i=1;i<=n;i++)
            g[i][j]=(g[i-1][j]+ksm(i-1,j))%mo;
    for(int j=1;j<19;j++)
        for(int i=1;i<=n;i++)
            f[i][j]=f[f[i][j-1]][j-1];
    int m=read();
    while(m--)
    {
        int x=read(),y=read(),k=read();
        int z=getlca(x,y);
        int ans=(g[dep[x]][k]+g[dep[y]][k])%mo;
        ans=(ans-(LL)2*g[dep[z]][k]%mo+mo)%mo;
        int sum=((LL)g[dep[z]][k]-g[dep[z]-1][k]+mo)%mo;
        ans=(ans+sum)%mo;
        write(ans),putchar('\n');
    }
    return 0;
}