【BZOJ 4034】【HAOI 2015】树上操作【树链剖分】

Description

有一棵点数为 N 的树，以点 1 为根，且树点有边权。然后有 M 个
操作，分为三种：
操作 1 ：把某个节点 x 的点权增加 a 。
操作 2 ：把某个节点 x 为根的子树中所有点的点权都增加 a 。
操作 3 ：询问某个节点 x 到根的路径中所有点的点权和。

Input

第一行包含两个整数 N, M 。表示点数和操作数。接下来一行 N 个整数，表示树中节点的初始权值。接下来 N-1 行每行三个正整数 fr, to ， 表示该树中存在一条边 (fr, to) 。再接下来 M 行，每行分别表示一次操作。其中第一个数表示该操作的种类（ 1-3 ） ，之后接这个操作的参数（ x 或者 x a ） 。

Output

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

题解

树链剖分。
值得注意是第二个操作，就是将tid[x]~tid[x]+size[x]-1的区间都加上a。
tid[]表示结点剖分后在线段树中的位置。
size[]表示以这个点为根的子树大小。
一定要注意开long long

代码

#include<cstdio>
#include<cstring>
#include<algorithm>

#define N 100010
typedef long long LL;
using namespace std;

struct node{int to,next;}e[N*2];
int head[N],m;
int dep[N],son[N],fa[N],top[N],size[N],tid[N],Rank[N];
int n,tim,a[N];

void init() {
    tim = m = 0;
    memset(head,0,sizeof(head));
    memset(son,-1,sizeof(son));
}

void add_edge(int from,int to) {
    e[++m].next = head[from];
    head[from] = m;
    e[m].to = to;
}

void dfs1(int v,int pa,int d)
{
    dep[v] = d; size[v] = 1;fa[v] = pa;
    for(int i = head[v];i;i=e[i].next)
        if(e[i].to != pa) {
            dfs1(e[i].to,v,d+1);
            size[v] += size[e[i].to];
            if(son[v] == -1 || size[e[i].to] > size[son[v]])
                son[v] = e[i].to;
        }
}

void dfs2(int v,int tp)
{
    top[v] = tp; tid[v] = ++tim;
    Rank[tid[v]] = v;
    if(son[v] == -1) return;
    dfs2(son[v],tp);
    for(int i = head[v];i;i=e[i].next)
        if(e[i].to != fa[v] && e[i].to != son[v])
            dfs2(e[i].to,e[i].to);
}

#define lson o << 1
#define rson o << 1 | 1
LL sum[N << 2],add[N << 2];

void build(int o,int l,int r)
{
    add[o] = 0;
    if(l == r){sum[o] = a[Rank[l]];return;}
    int mid = (l+r)>>1;
    build(lson,l,mid); build(rson,mid+1,r);
    sum[o] = sum[lson] + sum[rson];
}

void pushdown(int o,int l,int r)
{
    if(add[o] != 0) {
        int mid = (l+r)>>1;
        add[lson] += add[o]; sum[lson] += (LL)(mid-l+1)*add[o];
        add[rson] += add[o]; sum[rson] += (LL)(r-mid)*add[o];
        add[o] = 0;
    }
}

void update(int o,int l,int r,int pos,int v)
{
    if(l == r) {add[o] += v;sum[o] += v;return;}
    int mid = (l+r)>>1;
    pushdown(o,l,r);
    if(pos <= mid) update(lson,l,mid,pos,v); else update(rson,mid+1,r,pos,v);
    sum[o] = sum[lson] + sum[rson];
}

void change(int o,int l,int r,int ll,int rr,int v)
{
    if(ll <= l && rr >= r) {add[o] += v;sum[o] += (LL)(r-l+1)*v;return;}
    pushdown(o,l,r);
    int mid = (l+r)>>1;
    if(ll <= mid) change(lson,l,mid,ll,rr,v);
    if(rr > mid) change(rson,mid+1,r,ll,rr,v);
    sum[o] = sum[lson] + sum[rson];
}

LL query(int o,int l,int r,int ll,int rr)
{
    if(ll <= l && rr >= r) return sum[o];
    pushdown(o,l,r);
    int mid = (l+r)>>1;
    LL ans = 0;
    if(ll <= mid) ans += query(lson,l,mid,ll,rr);
    if(rr > mid) ans += query(rson,mid+1,r,ll,rr);
    return ans;
}

LL Query(int x,int y)
{
    LL ans = 0;
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        ans += query(1,1,n,tid[top[x]],tid[x]);
        x = fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x,y);
    ans += query(1,1,n,tid[x],tid[y]);
    return ans;
}

int main()
{
    int opt,u,v,q;
    scanf("%d%d",&n,&q);
    for(int i = 1;i <= n;i++) scanf("%d",&a[i]);
    init();
    for(int i = 1;i < n;i++) {
        scanf("%d%d",&u,&v);
        add_edge(u,v); add_edge(v,u);
    }
    dfs1(1,0,0);
    dfs2(1,1);
    build(1,1,n);
    while(q--) {
        scanf("%d",&opt);
        if(opt == 3) {
            scanf("%d",&u);
            printf("%lld\n",Query(1,u));
        }
        else {
            scanf("%d%d",&u,&v);
            if(opt == 1) update(1,1,n,tid[u],v);
            else change(1,1,n,tid[u],tid[u]+size[u]-1,v);
        }
    }
    return 0;
}