【DP】HDOJ 5426 Rikka with Game

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;

const int maxn = 100005;
const int maxm = 505;
const LL INF = 4e18;

struct Edge
{
	int v;
	Edge *next;
}*H[maxn], E[maxn << 1], *edges;

LL dp[maxn][maxm];
int size[maxn];
LL sum[maxn];
int a[maxn];
LL c[maxn];
int n, A, K;

void init()
{
	edges = E;
	for(int i = 0; i <= n; i++) {
		H[i] = 0;
		size[i] = 0;
		for(int j = 0; j <= K; j++)
			dp[i][j] = -INF;
	}
}

void addedges(int u, int v)
{
	edges->v = v;
	edges->next = H[u];
	H[u] = edges++;
}

void dfs(int u, int fa)
{
	sum[u] = a[u];
	for(Edge *e = H[u]; e; e = e->next) if(e->v != fa) {
		int v = e->v;
		dfs(v, u);
		sum[u] += sum[v];
	}

	dp[u][0] = a[u];
	size[u] = 1;
	for(Edge *e = H[u]; e; e = e->next) if(e->v != fa) {
		int v = e->v;
		for(int i = 0; i <= min(K, size[u] + size[v] + 1); i++) c[i] = -INF;
		for(int i = 0; i <= min(K, size[u]); i++) if(dp[u][i] != -INF) {
			c[i+1] = max(c[i+1], dp[u][i] + A * sum[v]);
			c[i] = max(c[i], dp[u][i]);
			for(int j = 0; j <= min(K, size[v]) && i + j <= K; j++) if(dp[v][j] != -INF)
				c[i + j] = max(c[i + j], dp[u][i] + dp[v][j]);
		}
		size[u] += size[v];
		for(int i = 0; i <= min(K, size[u]); i++) dp[u][i] = c[i];
	}
}

void work()
{
	for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
	for(int i = 1; i < n; i++) {
		int u, v;
		scanf("%d%d", &u, &v);
		addedges(u, v);
		addedges(v, u);
	}

	dfs(1, 1);
	LL ans = -INF;
	for(int i = 0; i <= K; i++) ans = max(ans, dp[1][i]);
	printf("%lld
", ans);
}

int main()
{
	while(scanf("%d%d%d", &n, &A, &K) != EOF) {
		init();
		work();
	}
	
	return 0;
}