bzoj5379 Tree dfs序列+线段树+lca

Description

---

给出一棵树，要求资磁一下三个操作
1. 换根为x
2. 将x的子树内所有点权加上y
3. 记x,y的lca为z，求z所在子树内所有点权之和

Solution

---

这道题一看就非常签到，实在不会还可以LCT水分

对于操作1我们不用真的换根，用一个变量root记录此时的根即可
对于操作2我们可以根据x和root的相对关系判断是dfs序上的连续一段还是它的补集
对于操作3我们可以发现一个结论：两点的lca一定是x、y、root两两lca中最深的那个

注意操作2可能会出现x=root的情况需要把整棵树都算一次

一开始写的是倍增求lca结果被捆绑测试卡掉17pts，后面改成重链剖分求lca就过了

Code

---

#include <stdio.h>
#include <string.h>
#include <algorithm>
#define rep(i,st,ed) for (int i=st;i<=ed;++i)
#define drp(i,st,ed) for (int i=st;i>=ed;--i)
#define fill(x,t) memset(x,t,sizeof(x))

typedef long long LL;
const int N=400005;
const int E=700005;

struct edge {int y,next;} e[E];

LL sum[N<<2],tag[N<<2];
int fa[N][19],pos[N],size[N],dep[N],bl[N];
int ls[N],w[N],edCnt,n,q;

int read() {
    int x=0,v=1; char ch=getchar();
    for (;ch<'0'||ch>'9';v=(ch=='-')?(-1):(v),ch=getchar());
    for (;ch<='9'&&ch>='0';x=x*10+ch-'0',ch=getchar());
    return x*v;
}

void add_edge(int x,int y) {
    e[++edCnt]=(edge) {y,ls[x]}; ls[x]=edCnt;
    e[++edCnt]=(edge) {x,ls[y]}; ls[y]=edCnt;
}

void push_down(int now,int l,int r,int mid) {
    if (!tag[now]) return ;
    LL w=tag[now]; tag[now]=0;
    tag[now<<1]+=w; sum[now<<1]+=(LL)(mid-l+1)*w;
    tag[now<<1|1]+=w; sum[now<<1|1]+=(LL)(r-mid)*w;
}

void modify(int now,int tl,int tr,int l,int r,LL v) {
    if (r<l) return ;
    if (tl==l&&tr==r) {
        sum[now]+=(LL)(r-l+1)*v;
        tag[now]+=v;
        return ;
    }
    int mid=(tl+tr)>>1;
    push_down(now,tl,tr,mid);
    if (r<=mid) modify(now<<1,tl,mid,l,r,v);
    else if (l>mid) modify(now<<1|1,mid+1,tr,l,r,v);
    else {
        modify(now<<1,tl,mid,l,mid,v);
        modify(now<<1|1,mid+1,tr,mid+1,r,v);
    }
    sum[now]=sum[now<<1]+sum[now<<1|1];
}

LL query(int now,int tl,int tr,int l,int r) {
    if (r<l) return 0;
    if (tl==l&&tr==r) return sum[now];
    int mid=(tl+tr)>>1;
    push_down(now,tl,tr,mid);
    if (r<=mid) return query(now<<1,tl,mid,l,r);
    if (l>mid) return query(now<<1|1,mid+1,tr,l,r);
    return query(now<<1,tl,mid,l,mid)+query(now<<1|1,mid+1,tr,mid+1,r);
}

void dfs1(int now) {
    rep(i,1,18) fa[now][i]=fa[fa[now][i-1]][i-1];
    size[now]=1; pos[now]=++pos[0];
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y==fa[now][0]) continue;
        fa[e[i].y][0]=now;
        dep[e[i].y]=dep[now]+1;
        dfs1(e[i].y); size[now]+=size[e[i].y];
    }
}

void dfs2(int now,int up) {
    bl[now]=up; int mx=0;
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y!=fa[now][0]&&size[e[i].y]>size[mx]) mx=e[i].y;
    }
    if (!mx) return ;
    dfs2(mx,up);
    for (int i=ls[now];i;i=e[i].next) {
        if (e[i].y!=fa[now][0]&&e[i].y!=mx) dfs2(e[i].y,e[i].y);
    }
}

int get_lca(int x,int y) {
    while (bl[x]!=bl[y]) {
        if (dep[bl[x]]<dep[bl[y]]) std:: swap(x,y);
        x=fa[bl[x]][0];
    }
    if (dep[x]<dep[y]) return x;
    return y;
}

int get_up(int st,int ed) {
    drp(i,18,0) if (dep[fa[st][i]]>dep[ed]) st=fa[st][i];
    return st;
}

void solve(int root,int x) {
    LL ans=0;
    if (x==root) ans=query(1,1,n,1,n);
    else if (get_lca(x,root)!=x) ans=query(1,1,n,pos[x],pos[x]+size[x]-1);
    else {
        int up=get_up(root,x);
        if (pos[up]>1) ans+=query(1,1,n,1,pos[up]-1);
        if (pos[up]+size[up]<=n) ans+=query(1,1,n,pos[up]+size[up],n);
    }
    printf("%lld\n", ans);
}

void change(int root,int x,LL v) {
    if (x==root) modify(1,1,n,1,n,v);
    else if (get_lca(x,root)!=x) modify(1,1,n,pos[x],pos[x]+size[x]-1,v);
    else {
        int up=get_up(root,x);
        if (pos[up]>1) modify(1,1,n,1,pos[up]-1,v);
        if (pos[up]+size[up]<=n) modify(1,1,n,pos[up]+size[up],n,v);
    }
}

int main(void) {
    freopen("tree.in","r",stdin);
    freopen("tree.out","w",stdout);
    n=read(),q=read(); int root=1;
    rep(i,1,n) w[i]=read();
    rep(i,2,n) add_edge(read(),read());
    dep[1]=1; dfs1(1); dfs2(1,1);
    rep(i,1,n) modify(1,1,n,pos[i],pos[i],w[i]);
    for (;q--;) {
        int opt=read(),x=read();
        if (opt==1) root=x;
        else if (opt==2) {
            int y=read(),v=read(),lca;
            int lx=get_lca(x,root),ly=get_lca(y,root),xy=get_lca(x,y);
            if (dep[lx]>=dep[ly]&&dep[lx]>=dep[xy]) lca=lx;
            if (dep[ly]>=dep[lx]&&dep[ly]>=dep[xy]) lca=ly;
            if (dep[xy]>=dep[lx]&&dep[xy]>=dep[ly]) lca=xy;
            change(root,lca,v);
        } else solve(root,x);
    }
    return 0;
}