BZOJ3252: 攻略

#include<iostream>
#include<cstdio>
#include<cstring>
#define N 200009
#define LL long long
using namespace std;

int n,k;

int sumedge,head[N],b[N];

int cnt,re[N],l[N],r[N],dad[N],vis[N];

LL ans,a[N],w[N];

inline int read(){
    char ch=getchar();int x=0,f=1;
    for(;!isdigit(ch);ch=getchar())if(ch==‘-‘)f=-1;
    for(;isdigit(ch);ch=getchar())x=x*10+ch-‘0‘;
    return x*f;
}

struct Edge{
    int x,y,nxt;
    Edge(int x=0,int y=0,int nxt=0):
        x(x),y(y),nxt(nxt){}
}edge[N];

struct Tree{
    int l,r,p;
    LL mx,s;
}tr[N<<2];

void add(int x,int y){
    edge[++sumedge]=Edge(x,y,head[x]);
    head[x]=sumedge;
}

void dfs(int x,LL sum){
    bool flag=true;l[x]=cnt+1;
    for(int i=head[x];i;i=edge[i].nxt){
        int v=edge[i].y;
        if(v==dad[x])continue;
        dad[v]=x;w[v]+=w[x];flag=false;
        dfs(v,w[v]);
    }
    if(flag)a[++cnt]=sum,re[cnt]=x;
    r[x]=cnt;
}

void pushup(int rt){
    tr[rt].mx=max(tr[rt<<1].mx,tr[rt<<1|1].mx);
    tr[rt].p=tr[rt<<1].mx>tr[rt<<1|1].mx?tr[rt<<1].p:tr[rt<<1|1].p;
    return;
}

void pushdown(int rt){
    if(tr[rt].s==0)return;
    tr[rt<<1].s+=tr[rt].s;tr[rt<<1|1].s+=tr[rt].s;
    tr[rt<<1].mx-=tr[rt].s;tr[rt<<1|1].mx-=tr[rt].s;
    tr[rt].s=0;
}

void build(int rt,int l,int r){
    tr[rt].l=l;tr[rt].r=r;
    if(l==r){
        tr[rt].mx=a[l];
        tr[rt].p=l;
        return;
    }
    int mid=(l+r)>>1;
    build(rt<<1,l,mid);build(rt<<1|1,mid+1,r);
    pushup(rt);
}

void change(int rt,int l,int r,int ql,int qr,int p){
    if(l>=ql&&r<=qr){
        tr[rt].mx-=p;
        tr[rt].s+=p;
        return;
    }
    pushdown(rt);
    int mid=(l+r)>>1;
    if(ql<=mid)change(rt<<1,l,mid,ql,qr,p);
    if(qr>mid)change(rt<<1|1,mid+1,r,ql,qr,p);
    pushup(rt);
}

int main(){
    n=read();k=read();
    for(int i=1;i<=n;i++)w[i]=read(),b[i]=w[i];
    for(int i=1;i<n;i++){
        int x,y;
        x=read();y=read(); 
        add(x,y);                              
    }
    dfs(1,w[1]);                               
    build(1,1,cnt);                             
    for(int i=1;i<=k;i++){
        LL t=tr[1].mx;int p=tr[1].p;          
        if(t<=0)break;ans+=t;                 
        p=re[p];                              
        while(p&&vis[p]==0){                   
            vis[p]=true;                       
            change(1,1,cnt,l[p],r[p],b[p]);    
            p=dad[p];                         
        }
    }
    printf("%lld\n",ans);
    return 0;
}