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,因为深度 di 是连续的,我们直接利用前缀和计算即可。
时间复杂度 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;
}