题解 追逐

传送门

原题

Code:

#include <bits/stdc++.h>
using namespace std;
#define INF 0x3f3f3f3f
#define N 100010
#define ll long long
// #define int long long

char buf[1<<21], *p1=buf, *p2=buf;
inline int read() {
	int ans=0, f=1; char c=getchar();
	while (!isdigit(c)) {if (c=='-') f=-f; c=getchar();}
	while (isdigit(c)) {ans=(ans<<3)+(ans<<1)+(c^48); c=getchar();}
	return ans*f;
}

int n, V;
int a[N], head[N], size;
int sta[N], top;
ll f[N][105], g[N][105], val[N], ans;
struct edge{int to, next;}e[N<<1];
inline void add(int s, int t) {e[++size].to=t; e[size].next=head[s]; head[s]=size;}

void dfs(int u, int fa) {
	// cout<<"dfs: "<<u<<' '<<fa<<endl;
	top=0;
	for (int i=0; i<V; ++i) g[u][i]=val[u]; g[u][V]=0;
	for (int i=head[u],v; ~i; i=e[i].next) {
		v = e[i].to;
		if (v==fa) continue;
		dfs(v, u);
		for (int j=0; j<=V; ++j) ans=max(ans, g[u][j]+f[v][j]);
		for (int j=0; j<=V; ++j) {
			g[u][j]=max(g[u][j], g[v][j]);
			if (j<V) g[u][j]=max(g[u][j], g[v][j+1]+val[u]-a[v]);
			f[u][j]=max(f[u][j], f[v][j]);
			if (j>0) f[u][j]=max(f[u][j], f[v][j-1]+val[u]-a[fa]);
		}
	}
	for (int i=head[u],v; ~i; i=e[i].next) {
		v = e[i].to;
		if (v!=fa) sta[++top]=v;
	}
	for (int i=0; i<V; ++i) g[u][i]=val[u]; g[u][V]=0;
	for (int i=0; i<=V; ++i) f[u][i]=0;
	while (top) {
		int v=sta[top--];
		for (int j=0; j<=V; ++j) ans=max(ans, g[u][j]+f[v][j]);
		for (int j=0; j<=V; ++j) {
			g[u][j]=max(g[u][j], g[v][j]);
			if (j<V) g[u][j]=max(g[u][j], g[v][j+1]+val[u]-a[v]);
			f[u][j]=max(f[u][j], f[v][j]);
			if (j>0) f[u][j]=max(f[u][j], f[v][j-1]+val[u]-a[fa]);
		}
	}
}

signed main()
{
	freopen("c.in", "r", stdin);
	freopen("c.out", "w", stdout);

	n=read(); V=read();
	memset(head, -1, sizeof(head));
	for (int i=1; i<=n; ++i) a[i]=read();
	for (int i=1,u,v; i<n; ++i) {
		u=read(); v=read();
		add(u, v); add(v, u);
		val[u]+=a[v]; val[v]+=a[u];
	}
	// cout<<"val: "; for (int i=1; i<=n; ++i) cout<<val[i]<<' '; cout<<endl;
	dfs(1, 0);
	printf("%lld\n", ans);
	// cout<<g[8][1]<<endl;

	return 0;
}