Bzoj4034:[HAOI2015]T2:树链剖分

题目链接：4034:[HAOI2015]T2

水树链剖分，注意开long long……

#include<cstdio>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#define LL long long
using namespace std;
const int maxn=300010;
int n,m,tot=1,h[maxn],pos[maxn];
struct edges{int to,next;}G[maxn*2];
int size[maxn],ind=0;
LL a[maxn],b[maxn],ret=0;
int dep[maxn],fa[maxn],Belong[maxn];
struct seg{
    int l,r; LL data,sum;
    seg *lc,*rc;
    seg():data(0),sum(0){}
};
seg *root=new seg();
bool vis[maxn];

void add(int x,int y){
    G[++tot].to=y; G[tot].next=h[x]; h[x]=tot;
}

void dfs1(int x,int deep){
    dep[x]=deep; size[x]=1; vis[x]=1;
    for (int i=h[x];i;i=G[i].next){
        int v=G[i].to;
        if (vis[v]) continue;
        fa[v]=x; dfs1(v,deep+1);
        size[x]+=size[v];
    }
}

void dfs2(int x,int L){
    int k=0; Belong[x]=L;
    pos[x]=++ind; b[ind]=a[x];
    for (int i=h[x];i;i=G[i].next)
       if (dep[x]<dep[G[i].to]&&size[G[i].to]>size[k])
          k=G[i].to;
    if (!k) return; dfs2(k,L);
    for (int i=h[x];i;i=G[i].next)
       if (dep[x]<dep[G[i].to]&&k!=G[i].to)
          dfs2(G[i].to,G[i].to);
}

void push_up(seg *p){
    if (p->l+1==p->r) return; p->sum=0;
    if (p->lc) p->sum+=p->lc->sum;
    if (p->rc) p->sum+=p->rc->sum;
}

void push_down(seg *p){
    if (p->l+1==p->r) return;
    if (!p->data) return;
    if (p->lc) p->lc->sum+=(p->lc->r-p->lc->l)*p->data,
               p->lc->data+=p->data;
    if (p->rc) p->rc->sum+=(p->rc->r-p->rc->l)*p->data;
               p->rc->data+=p->data;
    p->data=0;
}

void build(seg *p,int l,int r){
    p->l=l; p->r=r;
    if (l+1==r){
        p->sum=b[l]; p->data=0;
        p->lc=p->rc=NULL; return;
    }else if (l+1<r){
        int mid=(l+r)>>1;
        p->lc=new seg();
        p->rc=new seg();
        if (l<mid) build(p->lc,l,mid);
        else p->lc=NULL;
        if (mid<r) build(p->rc,mid,r);
        else p->rc=NULL;
        push_up(p);
    }
}

void update(seg *p,int l,int r,LL v){
    push_down(p);
    if (l<=p->l&&p->r<=r){
        p->sum+=1LL*(p->r-p->l)*v;
        p->data+=v; return;
    }else{
        int mid=(p->l+p->r)>>1;
        if (l<mid) update(p->lc,l,r,v);
        if (mid<r) update(p->rc,l,r,v);
        push_up(p);
    }
}

void ask(seg *p,int l,int r){
    push_down(p);
    if (l<=p->l&&p->r<=r){
        ret+=p->sum; return;
    }else{
        int mid=(p->l+p->r)>>1;
        if (l<mid) ask(p->lc,l,r);
        if (mid<r) ask(p->rc,l,r);
    }
}

void getans(int x,int y){
    LL sum=0;
    while (Belong[x]!=Belong[y]){
        ret=0; ask(root,pos[Belong[x]],pos[x]+1);
        sum+=ret; x=fa[Belong[x]];
    } ret=0; ask(root,pos[y],pos[x]+1);
    sum+=ret; printf("%lld\n",sum);
}

int main(){
    scanf("%d%d",&n,&m);
    for (int i=1;i<=n;++i) scanf("%lld",&a[i]);
    for (int i=1;i<n;++i){
        int x,y; scanf("%d%d",&x,&y);
        add(x,y); add(y,x);
    }
    dfs1(1,1); dfs2(1,1);
    build(root,1,ind+1);
    for (int i=1;i<=m;++i){
        int opt,x; LL y;
        scanf("%d",&opt);
        if (opt==1){
            scanf("%d%lld",&x,&y);
            update(root,pos[x],pos[x]+1,y);
        }else if (opt==2){
            scanf("%d%lld",&x,&y);
            update(root,pos[x],pos[x]+size[x],y);
        }else if (opt==3){
            scanf("%d",&x);
            getans(x,1);
        }
    }
}