BZOJ_3252_攻略_线段树+dfs序

#include <stdio.h>
#include <string.h>
#include <algorithm>
using namespace std;
#define N 200050
#define ls p<<1
#define rs p<<1|1
typedef long long ll;
int mx[N<<2];
ll t[N<<2],inc[N<<2],ans;
int head[N],to[N<<1],nxt[N<<1],val[N],cnt,kill[N],dfn[N],son[N],k,tot,fa[N],turn[N],n;
ll dis[N],dd[N];
inline void add(int u,int v) {
    to[++cnt]=v; nxt[cnt]=head[u]; head[u]=cnt;
}
void dfs(int x,int y) {
    dis[x]=dis[y]+val[x];
    int i,flg=0; fa[x]=y;
    dfn[x]=tot+1;
    for(i=head[x];i;i=nxt[i]) {
        if(to[i]!=y) flg=1;
    }
    if(!flg) {
        dfn[x]=son[x]=++tot;
        turn[tot]=x;
        return ;
    }
    dfn[x]=tot+1;
    for(i=head[x];i;i=nxt[i]) {
        if(to[i]!=y) {
            dfs(to[i],x);
        }
    }
    son[x]=tot;
}
void pushup(int p) {
    if(t[ls]>t[rs]) mx[p]=mx[ls],t[p]=t[ls];
    else mx[p]=mx[rs],t[p]=t[rs];
}
void build(int l,int r,int p) {
    if(l==r) {
        mx[p]=turn[l];
        t[p]=dis[turn[l]];
        return ;
    }
    int mid=(l+r)>>1;
    build(l,mid,ls); build(mid+1,r,rs);
    pushup(p);
}
void pushdown(int p) {
    ll d=inc[p];
    if(d) {
        t[ls]+=d; t[rs]+=d;
        inc[ls]+=d; inc[rs]+=d;
        inc[p]=0;
    }
}
void update(int l,int r,int x,int y,int v,int p) {
    if(x<=l&y>=r) {
        t[p]+=v; inc[p]+=v; return ;
    }
    pushdown(p);
    int mid=(l+r)>>1;
    if(x<=mid) update(l,mid,x,y,v,ls);
    if(y>mid) update(mid+1,r,x,y,v,rs);
    pushup(p);
}
void solve(int x) {
    if(kill[x]||!x) return ;
    kill[x]=1;
    update(1,tot,dfn[x],son[x],-val[x],1);
    solve(fa[x]);
}
int main() {
    scanf("%d%d",&n,&k);
    int i,x,y;
    for(i=1;i<=n;i++) {
        scanf("%d",&val[i]);
    }
    for(i=1;i<n;i++) {
        scanf("%d%d",&x,&y);
        add(x,y); add(y,x);
    }
    dfs(1,0);
    build(1,tot,1);
    for(i=1;i<=k;i++) {
        ans+=t[1];
        solve(mx[1]);
    }
    printf("%lld\n",ans);
}