[Bzoj3252]攻略（dfs序+线段树）

Description

题目链接

Solution

可以想到，每次肯定是拿最大价值为最优

考虑改变树上一个点的值，只会影响它的子树，也就是dfs序上的一个区间，

于是可以以dfs序建线段树，这样就变成区间问题了

Code

#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;
}