CF960H Santa's Gift

Description

给定一棵树，每个点初始一个颜色，每种颜色有一个权值 \(b_i\)。有两种操作，把一个点改成另一种颜色，或给定一个 \(k\)，询问下面式子的值。

\[\sum_{i=1}^n (S_ib_k-C)^2 \]

其中 \(S_i\) 表示子树 \(i\) 中颜色为 \(k\) 的点的个数，\(C\) 是一个常数，在题目开始时给出。

Solution

先化开这个式子，就是

\[b_k^2\sum S_i^2-2Cb_k \sum S_i+nC^2 \]

所以只需要维护 \(\sum S_i^2\) 和 \(\sum S_i\)，实际上不用每个点都维护一个 \(S_i^2\)，因为我们只需要知道和即可。所以可以考虑每次修改时计算 \(\sum S_i^2\) 的增量。假设当前修改的点是 \(i\)，那么平方和的增量就是

\[\sum_{u\in Path(1,i)} 2S_{u}+1=\Big(2\times \sum_{u\in Path(1,i)} S_{u} \Big)+ depth_{i} \]

一次和的增量就是

\[\sum_{u\in Path(1,i)} 1=depth_{i} \]

所以直接对每种颜色用树剖维护每个点的 \(S_i\)，注意要动态开点。

#include<stdio.h>
#include<algorithm>
using namespace std;

typedef long long ll;

const int N=5e4+7;

inline int read(){
    int x=0,flag=1; char c=getchar();
    while(c<'0'||c>'9'){if(c=='-')flag=0;c=getchar();}
    while(c>='0'&&c<='9'){x=(x<<1)+(x<<3)+c-48;c=getchar();}
    return flag? x:-x;
}

struct Node{
    int ls,rs;
    ll s,tag,len;
}t[N*160];

struct E{
    int next,to;
}e[N<<1];

int n,m,Q,head[N];
ll c[N],b[N],C,ans[N];
int dep[N],fa[N],sz[N],top[N],rt[N],son[N],in[N];

inline void add(int id,int to){
    static int cnt=0;
    e[++cnt]=(E){head[id],to};
    head[id]=cnt;
    e[++cnt]=(E){head[to],id};
    head[to]=cnt;
}

void dfs(int u){
    sz[u]=1;
    dep[u]=dep[fa[u]]+1;
    for(int i=head[u];i;i=e[i].next){
        int v=e[i].to;
        if(v==fa[u]) continue;
        fa[v]=u,dfs(v),sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}

void Dfs(int u,int tp){
    static int timer=0;
    in[u]=++timer,top[u]=tp;
    if(!son[u]) return ;
    Dfs(son[u],tp);
    for(int i=head[u];i;i=e[i].next){
        const int v=e[i].to;
        if(v!=son[u]&&v!=fa[u]) Dfs(v,v);
    }
}

inline void Add(int id,int v){
    t[id].s+=t[id].len*v;
    t[id].tag+=v;
}

int cnt=0;
inline void pd(int id){
    if(!t[id].tag) return ;
    Add(t[id].ls,t[id].tag);
    Add(t[id].rs,t[id].tag);
    t[id].tag=0;
}

int L,R,Val;
void modify(int &id,int lf,int rf){
    if(!id) id=++cnt,t[id].len=rf-lf+1;
    if(L<=lf&&rf<=R){Add(id,Val);return;}
    int mid=(lf+rf)>>1;
    if(!t[id].ls) t[id].ls=++cnt,t[cnt].len=mid-lf+1;
    if(!t[id].rs) t[id].rs=++cnt,t[cnt].len=rf-mid; pd(id);
    if(L<=mid) modify(t[id].ls,lf,mid);
    if(R>mid) modify(t[id].rs,mid+1,rf);
    t[id].s=t[t[id].ls].s+t[t[id].rs].s;
}

void Modify(int &Rt,int u,int v){
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        L=in[top[u]],R=in[u],modify(Rt,1,n);
        u=fa[top[u]];
    }
    if(dep[u]<dep[v]) swap(u,v);
    L=in[v],R=in[u],modify(Rt,1,n);
}

int query(int id,int lf,int rf){
    if(!id) return 0;
    if(L<=lf&&rf<=R) return t[id].s;
    int mid=(lf+rf)>>1,ret=0;
    if(!t[id].ls) t[id].ls=++cnt,t[cnt].len=mid-lf+1;
    if(!t[id].rs) t[id].rs=++cnt,t[cnt].len=rf-mid; pd(id);
    if(L<=mid) ret+=query(t[id].ls,lf,mid);
    if(R>mid) ret+=query(t[id].rs,mid+1,rf);
    return ret;
}

int Query(int Rt,int u,int v){
    int ret=0;
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        L=in[top[u]],R=in[u],ret+=query(Rt,1,n);
        u=fa[top[u]];
    }
    if(dep[u]<dep[v]) swap(u,v);
    L=in[v],R=in[u]; return ret+query(Rt,1,n);
}

void Add_(int i,int c){
    ans[c]=(ans[c]+1ll*(2*Query(rt[c],1,i)+dep[i])*b[c]*b[c]);
    ans[c]=(ans[c]-2ll*C*b[c]*dep[i]);
    Val=1,Modify(rt[c],1,i);
}

void Del_(int i,int c){
    Val=-1,Modify(rt[c],1,i);
    ans[c]=(ans[c]+2ll*C*b[c]*dep[i]);
    ans[c]=(ans[c]-1ll*(2*Query(rt[c],1,i)+dep[i])*b[c]*b[c]);
}

int main(){
    freopen("gift.in","r",stdin);
    freopen("gift.out","w",stdout);
    n=read(),m=read(),Q=read(),C=read();
    for(int i=1;i<=n;i++) c[i]=read();
    for(int i=2;i<=n;i++) add(i,read());
    for(int i=1;i<=m;i++)
        b[i]=read(),ans[i]=1ll*n*C*C;
    dfs(1),Dfs(1,1);
    for(int i=1;i<=n;i++) Add_(i,c[i]);
    while(Q--){
        int op=read();
        if(op==1){
            int p=read(),c_=read();
            if(c[p]==c_) continue;
            Del_(p,c[p]),Add_(p,c[p]=c_);
        }else printf("%lf\n",(double)ans[read()]/n);
    }
}

/*
4 3 3 10
1 1 2 1
12 14 7
2 3
3 4
1 4
2 1
1 1 2
2 2
*/