[树链剖分][树状数组] Luogu P3676 小清新数据结构题

题目描述

在很久很久以前，有一棵n个点的树，每个点有一个点权。

现在有q次操作，每次操作是修改一个点的点权或指定一个点，询问以这个点为根时每棵子树点权和的平方和。

（题目不是很好懂，没看太懂的可以看看样例解释）

题解

• 大爷博客（戳我）

代码

#include <cstdio>
#include <cstring>
#include <iostream> 
#define ll long long
#define sqr(x) (x)*(x)
using namespace std;
const ll N=2e5+10; 
struct edge{ll from,to;}e[N<<1];
ll n,m,cnt,ans,tot,w[N],c[N],fa[N],sz[N],id[N],pos[N],deep[N],son[N],top[N],head[N],sz1[N],sz2[N];
void add(ll x,ll y) { for (int i=x;i<=n;i+=i&(-i)) sz1[i]+=y,sz2[i]+=x*y; }
ll query(ll x) { ll r=0; for (ll i=x;i;i-=i&(-i)) r+=(x+1)*sz1[i]-sz2[i]; return r; }
void insert(ll x,ll y)
{
    e[++cnt].to=y,e[cnt].from=head[x],head[x]=cnt;
    e[++cnt].to=x,e[cnt].from=head[y],head[y]=cnt;
}
void dfs1(ll x)
{
    deep[x]=deep[fa[x]]+1,sz[x]=1;
    for (ll i=head[x];i;i=e[i].from) if (e[i].to!=fa[x])
    {
        fa[e[i].to]=x,dfs1(e[i].to),sz[x]+=sz[e[i].to],w[x]+=w[e[i].to];
        if (sz[e[i].to]>sz[son[x]]) son[x]=e[i].to;
    }
    ans+=sqr(w[x]);
}
void dfs2(ll x,ll pre)
{
    top[x]=pre;pos[id[x]=++tot]=x,add(id[x],w[x]-w[pos[tot-1]]);
    if (son[x]) dfs2(son[x],pre);
    for (ll i=head[x];i;i=e[i].from) if (fa[x]!=e[i].to&&e[i].to!=son[x]) dfs2(e[i].to,e[i].to);
}
void modify(ll x,ll y)
{
    ll l=0,r=0;
    for (;x;x=fa[top[x]]) l+=id[x]-id[top[x]]+1,r+=query(id[x])-query(id[top[x]]-1),add(id[top[x]],y),add(id[x]+1,-y);
    ans+=y*(y*l+(r<<1));
}
ll calc(ll u)
{
    ll k=0,x=query(1),s=0;
    for (;u;u=fa[top[u]]) k+=id[u]-id[top[u]]+1,s+=query(id[u])-query(id[top[u]]-1);
    return ans+x*((k+1)*x-(s<<1));
}
int main()
{
    scanf("%lld%lld",&n,&m);
    for (ll i=2,x,y;i<=n;i++) scanf("%lld%lld",&x,&y),insert(x,y);
    for (ll i=1;i<=n;i++) scanf("%lld",&c[i]),w[i]=c[i];
    dfs1(1),dfs2(1,1);
    for (ll op,x,y;m;m--)
    {
        scanf("%lld%lld",&op,&x);
        if (op==1) scanf("%lld",&y),y-=c[x],c[x]+=y,modify(x,y); else printf("%lld
",calc(x));
    }
}