【刷题】BZOJ 5293 [Bjoi2018]求和

Description

master 对树上的求和非常感兴趣。他生成了一棵有根树，并且希望多次询问这棵树上一段路径上所有节点深度的k

次方和，而且每次的k 可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给

了pupil，但pupil 并不会这么复杂的操作，你能帮他解决吗？

Input

第一行包含一个正整数n ，表示树的节点数。

之后n-1 行每行两个空格隔开的正整数i,j ，表示树上的一条连接点i 和点j 的边。

之后一行一个正整数m ，表示询问的数量。

之后每行三个空格隔开的正整数i,j,k ，表示询问从点i 到点j 的路径上所有节点深度的k 次方和。

由于这个结果可能非常大，输出其对998244353 取模的结果。

树的节点从1 开始标号，其中1 号节点为树的根。

Output

对于每组数据输出一行一个正整数表示取模后的结果。

1≤n,m≤300000,1≤k≤50

Sample Input

5

1 2

1 3

2 4

2 5

2

1 4 5

5 4 45

Sample Output

33

503245989

说明

样例解释

以下用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

这题就一水题，指数小于等于50，直接对于每一个指数都维护树上两点权值和

开始写了一个树剖，交洛谷加O2过了，然后发现跑得有点慢，就又写了一个差分

这是差分的代码：

#include<bits/stdc++.h>

#define ui unsigned int

#define ll long long

#define db double

#define ld long double

#define ull unsigned long long

const int MAXN=300000+10,Mod=998244353;

int n,q,e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],Jie[MAXN][20],dep[MAXN];

ll Sum[MAXN][51];

template<typename T> inline void read(T &x)

{

	T data=0,w=1;

	char ch=0;

	while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();

	if(ch=='-')w=-1,ch=getchar();

	while(ch>='0'&&ch<='9')data=((T)data<<3)+((T)data<<1)+(ch^'0'),ch=getchar();

	x=data*w;

}

template<typename T> inline void write(T x,char ch='\0')

{

	if(x<0)putchar('-'),x=-x;

	if(x>9)write(x/10);

	putchar(x%10+'0');

	if(ch!='\0')putchar(ch);

}

template<typename T> inline void chkmin(T &x,T y){x=(y<x?y:x);}

template<typename T> inline void chkmax(T &x,T y){x=(y>x?y:x);}

template<typename T> inline T min(T x,T y){return x<y?x:y;}

template<typename T> inline T max(T x,T y){return x>y?x:y;}

inline void insert(int x,int y)

{

	to[++e]=y;

	nex[e]=beg[x];

	beg[x]=e;

}

inline ll qexp(ll a,ll b)

{

	ll res=1;

	while(b)

	{

		if(b&1)res=res*a%Mod;

		a=a*a%Mod;

		b>>=1;

	}

	return res;

}

inline void dfs1(int x,int f)

{

	Jie[x][0]=f;dep[x]=dep[f]+1;

	for(register int i=beg[x];i;i=nex[i])

		if(to[i]==f)continue;

		else dfs1(to[i],x);

	for(register int i=1;i<=50;++i)Sum[x][i]=qexp(dep[x],i);

}

inline void dfs2(int x)

{

	for(register int i=1;i<=50;++i)(Sum[x][i]+=Sum[Jie[x][0]][i])%=Mod;

	for(register int i=beg[x];i;i=nex[i])

		if(to[i]==Jie[x][0])continue;

		else dfs2(to[i]);

}

inline void init()

{

	dfs1(1,0);

	dfs2(1);

	for(register int j=1;j<19;++j)

		for(register int i=1;i<=n;++i)Jie[i][j]=Jie[Jie[i][j-1]][j-1];

}

inline int LCA(int u,int v)

{

	if(dep[u]<dep[v])std::swap(u,v);

	if(dep[u]>dep[v])

		for(register int i=19;i>=0;--i)

			if(dep[Jie[u][i]]>=dep[v])u=Jie[u][i];

	if(u==v)return u;

	for(register int i=19;i>=0;--i)

		if(Jie[u][i]!=Jie[v][i])u=Jie[u][i],v=Jie[v][i];

	return Jie[u][0];

}

int main()

{

	read(n);

	for(register int i=1;i<n;++i)

	{

		int u,v;read(u);read(v);

		insert(u,v);insert(v,u);

	}

	dep[0]=-1;

	init();

	read(q);

	while(q--)

	{

		int u,v,k;read(u);read(v);read(k);

		int lca=LCA(u,v);

		write((Sum[u][k]+Sum[v][k]-Sum[lca][k]+Mod-Sum[Jie[lca][0]][k]+Mod)%Mod,'\n');

	}

	return 0;

}

这是树剖的代码（开O2能过，不知道不开O2或者其他OJ上能不能过）：

#include<bits/stdc++.h>

#define ui unsigned int

#define ll long long

#define db double

#define ld long double

#define ull unsigned long long

const int MAXN=300000+10,Mod=998244353;

int n,q,e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],st[MAXN],ed[MAXN],fa[MAXN],dep[MAXN],val[MAXN],hson[MAXN],size[MAXN],cnt,top[MAXN];

#define Mid ((l+r)>>1)

#define ls rt<<1

#define rs rt<<1|1

#define lson ls,l,Mid

#define rson rs,Mid+1,r

struct Segment_Tree{

    ll Sum[MAXN<<2][51];

    inline ll qexp(ll a,ll b)

    {

        ll res=1;

        while(b)

        {

            if(b&1)res=res*a%Mod;

            a=a*a%Mod;

            b>>=1;

        }

        return res;

    }

    inline void PushUp(int rt)

    {

        for(register int i=1;i<=50;++i)Sum[rt][i]=(Sum[ls][i]+Sum[rs][i])%Mod;

    }

    inline void Build(int rt,int l,int r)

    {

        if(l==r)

            for(register int i=1;i<=50;++i)Sum[rt][i]=qexp(val[l],i);

        else

        {

            Build(lson);Build(rson);

            PushUp(rt);

        }

    }

    inline ll Query(int rt,int l,int r,int L,int R,int k)

    {

        if(L<=l&&r<=R)return Sum[rt][k];

        else

        {

            ll res=0;

            if(L<=Mid)(res+=Query(lson,L,R,k))%=Mod;

            if(R>Mid)(res+=Query(rson,L,R,k))%=Mod;

            return res;

        }

    }

};

Segment_Tree T;

#undef Mid

#undef ls

#undef rs

#undef lson

#undef rson

template<typename T> inline void read(T &x)

{

    T data=0,w=1;

    char ch=0;

    while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();

    if(ch=='-')w=-1,ch=getchar();

    while(ch>='0'&&ch<='9')data=((T)data<<3)+((T)data<<1)+(ch^'0'),ch=getchar();

    x=data*w;

}

template<typename T> inline void write(T x,char ch='\0')

{

    if(x<0)putchar('-'),x=-x;

    if(x>9)write(x/10);

    putchar(x%10+'0');

    if(ch!='\0')putchar(ch);

}

template<typename T> inline void chkmin(T &x,T y){x=(y<x?y:x);}

template<typename T> inline void chkmax(T &x,T y){x=(y>x?y:x);}

template<typename T> inline T min(T x,T y){return x<y?x:y;}

template<typename T> inline T max(T x,T y){return x>y?x:y;}

inline void insert(int x,int y)

{

    to[++e]=y;

    nex[e]=beg[x];

    beg[x]=e;

}

inline void dfs1(int x,int f)

{

    int res=0;

    dep[x]=dep[f]+1;size[x]=1;fa[x]=f;

    for(register int i=beg[x];i;i=nex[i])

        if(to[i]==f)continue;

        else

        {

            dfs1(to[i],x);

            size[x]+=size[to[i]];

            if(size[to[i]]>res)res=size[to[i]],hson[x]=to[i];

        }

}

inline void dfs2(int x,int tp)

{

    top[x]=tp;st[x]=++cnt;val[cnt]=dep[x];

    if(hson[x])dfs2(hson[x],tp);

    for(register int i=beg[x];i;i=nex[i])

        if(to[i]==fa[x]||to[i]==hson[x])continue;

        else dfs2(to[i],to[i]);

    ed[x]=cnt;

}

inline void init()

{

    dfs1(1,0);

    dfs2(1,1);

    T.Build(1,1,n);

}

inline ll Getans(int u,int v,int k)

{

    ll res=0;

    while(top[u]!=top[v])

    {

        if(dep[top[u]]<dep[top[v]])std::swap(u,v);

        (res+=T.Query(1,1,n,st[top[u]],st[u],k))+=Mod;

        u=fa[top[u]];

    }

    (res+=T.Query(1,1,n,min(st[v],st[u]),max(st[v],st[u]),k))%=Mod;

    return res;

}

int main()

{

    read(n);

    for(register int i=1;i<n;++i)

    {

        int u,v;read(u);read(v);

        insert(u,v);insert(v,u);

    }

    dep[0]=-1;

    init();

    read(q);

    while(q--)

    {

        int u,v,k;read(u);read(v);read(k);

        write(Getans(u,v,k),'\n');

    }

    return 0;

}

巴特西