【Bzoj3252】攻略（dfs序+线段树）

#include <cstdio>
#include <algorithm>
#define MID int mid=(l+r)>>1,ls=id<<1,rs=id<<1|1
#define ll long long
#define N 200010
using namespace std;

struct xds{ll x,tag;}T[N<<2];
struct info{int to,nex;}e[N<<1];
int n,k,tot,head[N],val[N],dfn[N],bel[N],fa[N],R[N],x;
ll Ans,sum[N];
bool vis[N];

inline int read(){
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

void Link(int u,int v){
    e[++tot].nex=head[u];e[tot].to=v;head[u]=tot;
}

void dfs(int u){
    bel[dfn[u]=++tot]=u;sum[tot]=sum[dfn[fa[u]]]+val[u];
    for(int i=head[u];i;i=e[i].nex){
        int v=e[i].to;
        if(v==fa[u]) continue;
        fa[v]=u;
        dfs(v);
    }
    R[u]=tot;
}

void build(int l,int r,int id){
    if(l==r){T[id].x=sum[l];return;}
    MID;
    build(l,mid,ls);
    build(mid+1,r,rs);
    T[id].x=max(T[ls].x,T[rs].x);
}

void pushdown(int id){
    ll &tag=T[id].tag;
    if(!tag) return;
    int ls=id<<1,rs=id<<1|1;
    T[ls].x+=tag;T[ls].tag+=tag;
    T[rs].x+=tag;T[rs].tag+=tag;
    tag=0;
}

void Find(int l,int r,int id){
    if(l==r){x=bel[l];return;}
    pushdown(id);
    MID;
    if(T[ls].x>T[rs].x) Find(l,mid,ls);
    else Find(mid+1,r,rs);
    T[id].x=max(T[ls].x,T[rs].x);
}

void Modify(int l,int r,int id,int ql,int qr,int x){
    if(l>=ql&&qr>=r){
        T[id].x+=x;T[id].tag+=x;return;
    }
    pushdown(id);
    MID;
    if(ql<=mid) Modify(l,mid,ls,ql,qr,x);
    if(qr>mid) Modify(mid+1,r,rs,ql,qr,x);
    T[id].x=max(T[ls].x,T[rs].x);
}

int main(){
    n=read(),k=read();
    for(int i=1;i<=n;++i) val[i]=read();
    for(int i=1;i<n;++i){
        int u=read(),v=read();
        Link(u,v);Link(v,u);
    }
    tot=0;dfs(1);
    build(1,n,1);
    while(k--)
    {
        Ans+=T[1].x;
        Find(1,n,1);
        for(;x&&!vis[x];vis[x]=1,x=fa[x]) Modify(1,n,1,dfn[x],R[x],-val[x]);
    }
    printf("%lld\n",Ans);
    return 0;
}