2022HDU-Multi-University Training Rock Tree

题目链接：https://acm.hdu.edu.cn/contest/problem?cid=1050&pid=1011

大意就是求一个权值和最大的连通块，其直径不超过k。

似乎是个很常见的经典题，dp状态和深度有关，用长链剖分合并一下。

#include<bits/stdc++.h>
#define N 100009
using namespace std;
typedef long long ll;
int head[N],tot,dfn[N],tp[N];
int n,k,son[N],dep[N];
ll tr[N<<2],la[N<<2],a[N],ans,b[N],c[N]; 
inline ll rd(){
	ll x=0;scanf("%lld",&x);
	return x;
}
struct edge{
	int n,to;
}e[N<<1];
inline void pushdown(int cnt){
	la[cnt<<1]+=la[cnt];
	la[cnt<<1|1]+=la[cnt];
	tr[cnt<<1]+=la[cnt];
	tr[cnt<<1|1]+=la[cnt];
	la[cnt]=0;
}
void upd(int cnt,int l,int r,int L,int R,ll x){
	if(l>=L&&r<=R){
		tr[cnt]+=x;
		la[cnt]+=x;
		return;
	}
	int mid=(l+r)>>1;
	if(la[cnt])pushdown(cnt);
	if(mid>=L)upd(cnt<<1,l,mid,L,R,x);
	if(mid<R)upd(cnt<<1|1,mid+1,r,L,R,x);
	tr[cnt]=max(tr[cnt<<1],tr[cnt<<1|1]);
}
ll query(int cnt,int l,int r,int L,int R){
	if(l>=L&&r<=R)return tr[cnt];
	int mid=(l+r)>>1;ll ans=-1e18;
	if(la[cnt])pushdown(cnt);
	if(mid>=L)ans=max(ans,query(cnt<<1,l,mid,L,R));
	if(mid<R)ans=max(ans,query(cnt<<1|1,mid+1,r,L,R));
	return ans;
}
void updmax(int cnt,int l,int r,int x,ll y){
	if(l==r){
		tr[cnt]=max(tr[cnt],y);
		return;
	}
	int mid=(l+r)>>1;
	if(la[cnt])pushdown(cnt); 
	if(mid>=x)updmax(cnt<<1,l,mid,x,y);
	else updmax(cnt<<1|1,mid+1,r,x,y);
	tr[cnt]=max(tr[cnt<<1],tr[cnt<<1|1]);
}
inline void add(int u,int v){
	e[++tot].n=head[u];
	e[tot].to=v;
	head[u]=tot;
}
void dfs(int u,int fa){
	for(int i=head[u];i;i=e[i].n)if(e[i].to!=fa){
		int v=e[i].to;
		dfs(v,u);
		dep[u]=max(dep[v]+1,dep[u]);
		if(dep[v]>=dep[son[u]])son[u]=v;
	}
}
void dfs2(int u,int fa){
	dfn[u]=++dfn[0];
	if(son[u])dfs2(son[u],u),tp[u]=tp[son[u]];
	else tp[u]=u;
	upd(1,1,n,dfn[u],dfn[tp[u]],a[u]);
	for(int i=head[u];i;i=e[i].n)if(e[i].to!=fa){
		int v=e[i].to;
		if(v!=son[u]){
			dfs2(v,u);
			int md=min(k-1,dfn[tp[v]]-dfn[v]+1);
			c[0]=a[u]; 
			for(int j=1;j<=md;++j){
				b[j]=query(1,1,n,dfn[v]+j-1,dfn[v]+j-1);
				c[j]=query(1,1,n,dfn[u]+j,dfn[u]+j);
				c[j]=max(c[j],c[j-1]);
				b[j]=max(b[j],b[j-1]);
				updmax(1,1,n,dfn[u]+j,max(b[j]+c[min(k-j-1,j)],b[min(k-j-1,j)]+c[j]));
			}
			int l=md+2;
			int dd=dfn[tp[u]]-dfn[u]+1;
			for(int j=md;j>=1;--j){
				int r=min(k-j,dd);
				if(l<=r){
					upd(1,1,n,dfn[u]+l-1,dfn[u]+r-1,b[j]);
					l=r+1;
				}
			}
		}
	}
	ans=max(ans,query(1,1,n,dfn[u],min(dfn[u]+k-1,dfn[tp[u]])));
}
inline void sol(){
	n=rd();k=rd();
	ans=-1e18;
	for(int i=1;i<=n;++i)a[i]=rd();
	int u,v;
	for(int i=1;i<n;++i){
		u=rd();v=rd();
		add(u,v);add(v,u);
	}
	dfs(1,0);
	dfs2(1,0);
	printf("%lld\n",ans);
	for(int i=1;i<=n*4;++i)tr[i]=la[i]=0;
	for(int i=0;i<=n;++i){
		dfn[i]=tp[i]=head[i]=son[i]=dep[i]=0;
	}
	tot=0;
}
int main(){
//	freopen("1011.in","r",stdin);
	int T=rd();
	while(T--){
		sol();
	}
    return 0;
}