并不对劲的p3676:小清新数据结构题

题目大意

有一棵有\(n\)(\(n\leq 2*10^5\))个点的树，要进行\(q\)(\(q\leq 2*10^5\))次操作，每次操作是以下两种中的一种：

1.修改一个点的点权

2.指定一个点\(x\)，将该点变成根，并询问此时所有点的子树点权和之平方和

题解

设\(w_i\)表示以1号点为根时，点\(i\)的子树点权和

1操作可以看成将点\(x\)到1号点的路径上的\(w\)都加上\(y\)

假设点\(x\)到\(1\)的路径上的点是\(a_1,a_2,...a_b\),那么所有\(w\)的平方和增加了\((w_{a_1}+y)^2+(w_{a_2}+y)^2+...(w_{a_b}+y)^2-(w_{a_1}^2+w_{a_2}^2+...+w{a_b}^2)\)

该式=\(w_{a_1}^2+2*w_{a_1}*y+y^2+w_{a_2}^2+2*w_{a_2}*y+y^2+...+w_{a_b}^2+2*w_{a_b}*y+y^2-(w_{a_1}^2+w_{a_2}^2+...+w{a_b}^2)\)

=\((w_{a_1}^2+w_{a_2}^2+...+w_{a_b}^2)+2*y*(w_{a_1}+w_{a_2}+...+w_{a_b})+b*y^2-(w_{a_1}^2+w_{a_2}^2+...+w{a_b}^2)\)

=\(2*y*(w_{a_1}+w_{a_2}+...+w_{a_b})+b*y^2\)

2操作已知所有\(w\)的平方和，考虑把根从1号点挪到\(x\)，假设点\(x\)到\(1\)的路径上的点是\(a_1,a_2,...a_b\)，其中\(a_1\)为\(x\)，\(a_b\)为\(1\)

发现只有点\(x\)到\(1\)的路径上的点的子树点权和不是\(w\)

\(a_i(i\in[2,b])\)的子树点权和是\(w_{a_b}-w_{a_{i-1}}\)

\(a_1\)的子树点权和是\(w_{a_b}\)

那么 所有点的子树点权和之平方和 比 所有\(w\)的平方和 多出了 \(w_{a_b}^2+(w_{a_b}-w_{a_{1}})^2+(w_{a_b}-w_{a_{2}})^2+...+(w_{a_b}-w_{a_{b-1}})^2-(w_{a_1}^2+w_{a_2}^2+...+w_{a_b}^2)\)

该式=\(w_{a_b}^2+w_{a_b}^2-2*w_{a_b}*w_{a_1}+w_{a_{1}}^2+w_{a_b}^2-2*w_{a_b}*w_{a_2}+w_{a_{2}}^2+...+w_{a_b}^2-2*w_{a_b}*w_{a_{b-1}}+w_{a_{b-1}}^2-(w_{a_1}^2+w_{a_2}^2+...+w_{a_b}^2)\)

=\((b-1)*w_{a_b}^2-2*w_{a_b}*(w_{a_1}+...+w_{a_{b-1}})+(w_{a_1}^2+w_{a_2}^2+...+w_{a_b}^2)-(w_{a_1}^2+w_{a_2}^2+...+w_{a_b}^2)\)

=\((b-1)*w_{a_b}^2-2*w_{a_b}*(w_{a_1}+...+w_{a_{b-1}})\)

代码

#include<algorithm>
#include<cmath>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<ctime>
#include<iomanip>
#include<iostream>
#include<map>
#include<queue>
#include<set>
#include<stack>
#include<vector>
#define rep(i,x,y) for(register int i=(x);i<=(y);++i)
#define dwn(i,x,y) for(register int i=(x);i>=(y);--i)
#define maxn 200010
#define maxm (maxn<<1)
#define LL long long
#define ls (u<<1)
#define rs (u<<1|1)
#define mi (l+r>>1)
#define view(u,k) for(int k=fir[u];k!=-1;k=nxt[k])
using namespace std;
int read()
{
    int x=0,f=1;char ch=getchar();
    while(!isdigit(ch)&&ch!='-')ch=getchar();
    if(ch=='-')f=-1,ch=getchar();
    while(isdigit(ch))x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
    return x*f;
}
void write(LL x)
{
    if(x==0){putchar('0'),putchar('\n');return;}
    int f=0;char ch[20];
    if(x<0)putchar('-'),x=-x;
    while(x)ch[++f]=x%10+'0',x/=10;
    while(f)putchar(ch[f--]);
    putchar('\n');
    return;
}
int n,q,fir[maxn],nxt[maxm],vv[maxm],dfn[maxn],top[maxn],fa[maxn],tim,dep[maxn],siz[maxn],cnt;
LL key[maxn],sum[maxn],sumall,son[maxn],fakeans,mk[maxn<<2],tr[maxn<<2],to[maxn];
void ade(int u1,int v1){vv[cnt]=v1,nxt[cnt]=fir[u1],fir[u1]=cnt++;}
void getson(int u)
{
    siz[u]=1,sum[u]=key[u];
    view(u,k)
    {
        int v=vv[k];
        if(v==fa[u])continue;
        fa[v]=u,dep[v]=dep[u]+1,getson(v),siz[u]+=siz[v],sum[u]+=sum[v];
        if(siz[v]>siz[son[u]])son[u]=v;
    }fakeans+=sum[u]*sum[u];return;
}
void gettop(int u,int anc)
{
    dfn[u]=++tim,top[u]=anc,to[tim]=u;
    if(son[u])gettop(son[u],anc);
    view(u,k){int v=vv[k];if(v!=fa[u]&&v!=son[u])gettop(v,v);}
    return;
}
void pup(int u){tr[u]=tr[ls]+tr[rs];return;}
void mkad(int u,int l,int r,int k){tr[u]+=(LL)k*(LL)(r-l+1);mk[u]+=k;return;}
void pdn(int u,int l,int r){if(mk[u])mkad(ls,l,mi,mk[u]),mkad(rs,mi+1,r,mk[u]),mk[u]=0;}
void build(int u,int l,int r)
{
    if(l==r){tr[u]=sum[to[l]];return;}
    build(ls,l,mi),build(rs,mi+1,r),pup(u);
}
void add(int u,int l,int r,int x,int y,int k)
{
    if(x<=l&&r<=y)return mkad(u,l,r,k);
    pdn(u,l,r);
    if(x<=mi)add(ls,l,mi,x,y,k);
    if(y>mi)add(rs,mi+1,r,x,y,k);
    return pup(u);
}
LL ask(int u,int l,int r,int x,int y)
{
    if(x<=l&&r<=y)return tr[u];
    pdn(u,l,r);
    LL res=0;
    if(x<=mi)res+=ask(ls,l,mi,x,y);
    if(y>mi)res+=ask(rs,mi+1,r,x,y);
    return res;
}
LL askrd(int u)
{
    LL res=0;
    while(top[u]!=1)res+=ask(1,1,n,dfn[top[u]],dfn[u]),u=fa[top[u]];
    res+=ask(1,1,n,1,dfn[u]);
    return res;
}
LL askrd2(int u)
{
    LL res=0;
    while(top[u]!=1)res+=ask(1,1,n,dfn[top[u]],dfn[u]),u=fa[top[u]];
    if(u!=1)res+=ask(1,1,n,2,dfn[u]);
    return res;
}
void addrd(int u,int k)
{
    while(top[u]!=1)add(1,1,n,dfn[top[u]],dfn[u],k),u=fa[top[u]];
    add(1,1,n,1,dfn[u],k);return;
}
int main()
{
    n=read(),q=read();
    memset(fir,-1,sizeof(fir));
    rep(i,1,n-1){int x=read(),y=read();ade(x,y),ade(y,x);}
    rep(i,1,n)key[i]=sum[i]=read();
    getson(1),gettop(1,1),build(1,1,n);
    while(q--)
    {
        int f=read();
        if(f==1){int x=read();LL y=read();y-=key[x],key[x]+=y,fakeans+=2ll*y*askrd(x)+y*y*(dep[x]+1),addrd(x,y);}
        else{int x=read();LL tmp=ask(1,1,n,1,1);write(fakeans+dep[x]*tmp*tmp-2ll*tmp*askrd2(x));}
    }
    return 0;
}