【刷题】洛谷 P3676 小清新数据结构题

#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=200000+10,inf=0x3f3f3f3f;
int n,q,e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],sum[MAXN],vl[MAXN],st[MAXN],ed[MAXN],cnt,fa[MAXN],top[MAXN],hson[MAXN],size[MAXN],val[MAXN],dep[MAXN];
ll stans;
#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 Sum1[MAXN<<2],Sum2[MAXN<<2],Add[MAXN<<2];
    inline void PushUp(int rt)
    {
        Sum1[rt]=Sum1[ls]+Sum1[rs];
        Sum2[rt]=Sum2[ls]+Sum2[rs];
    }
    inline void PushDown(int rt,int len)
    {
        Sum2[ls]+=1ll*Add[rt]*Add[rt]*(len-(len>>1))+2ll*Add[rt]*Sum1[ls];
        Sum2[rs]+=1ll*Add[rt]*Add[rt]*(len>>1)+2ll*Add[rt]*Sum1[rs];
        Sum1[ls]+=1ll*Add[rt]*(len-(len>>1));
        Sum1[rs]+=1ll*Add[rt]*(len>>1);
        if(Add[ls]==inf)Add[ls]=0;
        if(Add[rs]==inf)Add[rs]=0;
        Add[ls]+=Add[rt];Add[rs]+=Add[rt];
        Add[rt]=inf;
    }
    inline void Build(int rt,int l,int r)
    {
        Add[rt]=inf;
        if(l==r)
        {
            Sum1[rt]=vl[l];
            Sum2[rt]=1ll*vl[l]*vl[l];
        }
        else
        {
            Build(lson);Build(rson);
            PushUp(rt);
        }
    }
    inline void Update(int rt,int l,int r,int L,int R,ll k)
    {
        if(L<=l&&r<=R)
        {
            Sum2[rt]+=1ll*k*k*(r-l+1)+2ll*k*Sum1[rt];
            Sum1[rt]+=1ll*(r-l+1)*k;
            if(Add[rt]==inf)Add[rt]=0;
            Add[rt]+=k;
        }
        else
        {
            if(Add[rt]!=inf)PushDown(rt,r-l+1);
            if(L<=Mid)Update(lson,L,R,k);
            if(R>Mid)Update(rson,L,R,k);
            PushUp(rt);
        }
    }
    inline ll Query(int rt,int l,int r,int L,int R)
    {
        if(L<=l&&r<=R)return Sum1[rt];
        else
        {
            ll res=0;
            if(Add[rt]!=inf)PushDown(rt,r-l+1);
            if(L<=Mid)res+=Query(lson,L,R);
            if(R>Mid)res+=Query(rson,L,R);
            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 dfs(int x,int f)
{
    sum[x]=val[x];
    for(register int i=beg[x];i;i=nex[i])
        if(to[i]==f)continue;
        else dfs(to[i],x),sum[x]+=sum[to[i]];
}
inline void dfs1(int x,int f)
{
    int res=0;
    size[x]=1;fa[x]=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);
            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)
{
    st[x]=++cnt;vl[cnt]=sum[x];top[x]=tp;
    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);stans=T.Sum2[1];
}
inline void Doans(int x,ll k)
{
    while(x)
    {
        T.Update(1,1,n,st[top[x]],st[x],k);
        x=fa[top[x]];
    }
}
inline ll Getans(int x)
{
    if(x==1)return stans;
    ll a0=T.Query(1,1,n,st[1],st[1]),res=stans+1ll*(dep[x]+1)*a0*a0,nsum=0;
    while(x)
    {
        nsum+=T.Query(1,1,n,st[top[x]],st[x]);
        x=fa[top[x]];
    }
    return res-2ll*a0*nsum;
}
int main()
{
    read(n);read(q);
    for(register int i=1;i<n;++i)
    {
        int u,v;read(u);read(v);
        insert(u,v);insert(v,u);
    }
    for(register int i=1;i<=n;++i)read(val[i]);
    dfs(1,0);
    init();
    while(q--)
    {
        int opt;read(opt);
        if(opt==1)
        {
            int x,y;read(x);read(y);
            Doans(x,y-val[x]);
            val[x]=y;stans=T.Sum2[1];
        }
        if(opt==2)
        {
            int x;read(x);
            write(Getans(x),'\n');
        }
    }
    return 0;
}