BZOJ3252: 攻略

#include<cstdio>
#include<cctype>
#include<queue>
#include<cstring>
#include<algorithm>
#define rep(i,s,t) for(int i=s;i<=t;i++)
#define dwn(i,s,t) for(int i=s;i>=t;i--)
#define ren for(int i=first[x];i;i=next[i])
using namespace std;
const int BufferSize=1<<16;
char buffer[BufferSize],*head,*tail;
inline char Getchar() {
	if(head==tail) {
		int l=fread(buffer,1,BufferSize,stdin);
		tail=(head=buffer)+l;
	}
	return *head++;
}
inline int read() {
    int x=0,f=1;char c=Getchar();
    for(;!isdigit(c);c=Getchar()) if(c==‘-‘) f=-1;
    for(;isdigit(c);c=Getchar()) x=x*10+c-‘0‘;
    return x*f;
}
typedef long long ll;
const int maxn=200010;
int n,k,A[maxn],first[maxn],next[maxn],to[maxn],e;
void AddEdge(int v,int u) {to[++e]=v;next[e]=first[u];first[u]=e;}
int st[maxn],en[maxn],pos[maxn],fa[maxn],del[maxn],cnt;
ll val[maxn];
void dfs(int x) {
	st[x]=++cnt;pos[cnt]=x;val[x]+=A[x];
	ren val[to[i]]=val[x],fa[to[i]]=x,dfs(to[i]);
	en[x]=cnt;
}
struct Ans {int p;ll v;}maxv[maxn<<2];
Ans better(Ans A,Ans B) {
	if(A.v>B.v) return A;
	return B;
}
ll addv[maxn<<2];
void maintain(int o,int l,int r) {
	int lc=o<<1,rc=lc|1;
	if(l<r) maxv[o]=better(maxv[lc],maxv[rc]);
	else maxv[o]=(Ans){l,0};
	if(addv[o]) maxv[o].v+=addv[o];
}
void build(int o,int l,int r) {
	if(l==r) addv[o]=val[pos[l]];
	else {
		int mid=l+r>>1,lc=o<<1,rc=lc|1;
		build(lc,l,mid);build(rc,mid+1,r);
	}
	maintain(o,l,r);
}
void update(int o,int l,int r,int ql,int qr,int v) {
	if(ql<=l&&r<=qr) addv[o]+=v;
	else {
		int mid=l+r>>1,lc=o<<1,rc=lc|1;
		if(addv[o]) addv[lc]+=addv[o],addv[rc]+=addv[o],addv[o]=0;
		if(ql<=mid) update(lc,l,mid,ql,qr,v);
		else maintain(lc,l,mid);
		if(qr>mid) update(rc,mid+1,r,ql,qr,v);
		else maintain(rc,mid+1,r);
	}
	maintain(o,l,r);
}
int main() {
	n=read();k=read();
	rep(i,1,n) A[i]=read();
	rep(i,1,n-1) AddEdge(read(),read());
	dfs(1);build(1,1,n);ll ans=0;
	while(k--) {
		Ans x=maxv[1];ans+=x.v;
		if(!x.v) break;
		int u=pos[x.p];
		while(u&&!del[u]) {
			del[u]=1;
			update(1,1,n,st[u],en[u],-A[u]);
			u=fa[u];
		}
	}
	printf("%lld\n",ans);
	return 0;
}