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

题目大意

有一棵有(n)((nleq 2*10^5))个点的树，要进行(q)((qleq 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(iin[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('
');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('
');
    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;
}