cogs 1963. [HAOI 2015] 树上操作 树链剖分+线段树

1963. [HAOI 2015] 树上操作

★★★☆   输入文件：haoi2015_t2.in   输出文件：haoi2015_t2.out   简单对比
时间限制：1 s   内存限制：256 MB

【题目描述】

有一棵点数为N的树，以点1为根，且树点有权值。然后有M个操作，分为三种：

操作1：把某个节点x的点权增加a。

操作2：把某个节点x为根的子树中所有点的点权都增加a。

操作3：询问某个节点x到根的路径中所有点的点权和。

【输入格式】

第一行两个整数N,M，表示点数和操作数。

接下来一行N个整数，表示树中节点的初始权值。

接下来N-1行每行两个正整数fr,to，表示该树中存在一条边(fr,to)。

再接下来M行，每行分别表示一次操作。其中第一个数表示该操作的种类（1~3），之后接这个操作的参数（x或者x a）。

【输出格式】

对于每个询问操作，输出该询问的答案。答案之间用换行隔开。

【样例输入】

5 5

1 2 3 4 5

1 2

1 4

2 3

2 5

3 3

1 2 1

3 5

2 1 2

3 3

【样例输出】

6

9

13

【提示】

对于30%的数据，N,M<=1000

对于50%的数据，N,M<=100000且数据随机。

对于100%的数据，N,M<=100000，且所有输入数据的绝对值都不会超过10^6。

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<string>
#include<cstring>
#include<cmath>
#include<vector>
using namespace std;
#define LL long long
const int maxn=100005;
int n,m,cot=0;
LL a[maxn];
vector<int> v[maxn];
int size[maxn];
int son[maxn];
int fa[maxn];
int top[maxn];
int dfn[maxn];
int pos[maxn];
int dep[maxn];
int en[maxn];
LL ls[maxn<<1],rs[maxn<<1],lz[maxn<<1],sum[maxn<<1];
int cnt=0;
void Dfs(int rt){
    size[rt]=1;
    for(int i=0;i<v[rt].size();i++)
        if(!size[v[rt][i]]){
            int to=v[rt][i];
            dep[to]=dep[rt]+1;
            fa[to]=rt;
            Dfs(to);
            size[rt]+=size[to];
            if(size[son[rt]]<size[to]) son[rt]=to;
        }
}
void Dfs(int rt,int tp){
    top[rt]=tp;
    dfn[++cnt]=rt;
    pos[rt]=cnt;
    if(son[rt]) Dfs(son[rt],tp);
    for(int i=0;i<v[rt].size();i++)
        if(!top[v[rt][i]])
            Dfs(v[rt][i],v[rt][i]);
    en[rt]=cnt;
}
int Build(int l,int r){
    int rt=++cot;
    if(l==r){
        sum[rt]=a[dfn[l]];
        return rt;
    }
    int mid=(l+r)>>1;
    ls[rt]=Build(l,mid);
    rs[rt]=Build(mid+1,r);
    sum[rt]=sum[ls[rt]]+sum[rs[rt]];
    return rt;
}
void Push_down(int rt,int l,int r){
    int mid=(l+r)>>1;
    lz[ls[rt]]+=lz[rt];
    lz[rs[rt]]+=lz[rt];
    sum[ls[rt]]+=lz[rt]*(mid-l+1);
    sum[rs[rt]]+=lz[rt]*(r-mid);
    lz[rt]=0;
    
}
LL Add(int rt,int l,int r,int s,int t,int qx){
    if(s>r||t<l)     return 0;
    if(s<=l&&r<=t){
        sum[rt]+=qx*1ll*(r-l+1);
        lz[rt]+=qx;
        return 0;
    }
    Push_down(rt,l,r);
    int mid=(l+r)>>1;
    Add(ls[rt],l,mid,s,t,qx);Add(rs[rt],mid+1,r,s,t,qx);
    sum[rt]=sum[ls[rt]]+sum[rs[rt]];
}
LL Sum(int rt,int l,int r,int s,int t){
    if(s>r||t<l) return 0;
    if(s<=l&&r<=t) return sum[rt];
    Push_down(rt,l,r);
    int mid=(l+r)>>1;
    return Sum(ls[rt],l,mid,s,t)+Sum(rs[rt],mid+1,r,s,t);
}
LL LCA_dis(int x,int y){
    LL res=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        res+=Sum(1,1,n,pos[top[x]],pos[x]);
        x=fa[top[x]];
    }
    if(dep[x]>dep[y])swap(x,y);
    res+=Sum(1,1,n,pos[x],pos[y]);
    return res;
}
int main()
{
    freopen("haoi2015_t2.in","r",stdin);
    freopen("haoi2015_t2.out","w",stdout);
    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);
        v[x].push_back(y);v[y].push_back(x);
    }
    Dfs(1);Dfs(1,1);Build(1,n);
    while(m--){
        int opt;scanf("%d",&opt);
        if(opt==1){
            int x,aa;scanf("%d%d",&x,&aa);
            Add(1,1,n,pos[x],pos[x],aa);
        }
        if(opt==2){
            int x,aa;
            scanf("%d%d",&x,&aa);
            Add(1,1,n,pos[x],en[x],aa);
        }
        if(opt==3){
            int x;scanf("%d",&x);
            printf("%lld
",LCA_dis(x,1));
        }
    }
    return 0;
}