问题 B: 【TEST 3】旅行

分析

此题与书上01背包有相似之处，观察题目，发现状态十分好设：第一维为节点，第二维为旅游的城市。于是，我们可以得到方程 \(dp_{i,j}=\max\limits_{k \in i} \{ \max \{ dp_{k,j},dp_{k,j-1}+a_i \} \}\)，可以打出代码

代码

#include <bits/stdc++.h>
#define int long long // 注意，开long long

using namespace std;

struct FastIO {
	template <typename T> FastIO& operator >> (T& in) {
		in = 0;
		char ch = getchar ();
		int flag = 1;
		for (; ! isdigit (ch); ch = getchar ()) if (ch == ‘-‘) flag = -1;
		for (;   isdigit (ch); ch = getchar ()) in = (in << 3) + (in << 1) + (ch ^ 48);
		in *= flag;
		return *this;
	}
}fin;

const int MaxN = 1e4 + 100, MaxK = 1e3 + 100;
int n, k;
int a[MaxN];
int nxt[MaxN << 1], to[MaxN << 1], head[MaxN], cnt;
int dp[MaxN][MaxK];

void add (int u, int v) {
	++ cnt;
	to[cnt] = v;
	nxt[cnt] = head[u];
	head[u] = cnt;
	return ;
}

void dfs (int d, int pre) { // 树形dp递归板子
	bool flag = 1;
	for (int i = head[d]; i; i = nxt[i]) {
		if (pre != to[i]) {
			flag = 0;
			dfs (to[i], d);
			for (int j = 0; j <= k; ++ j) {
				if (j == 0) {
					dp[d][j] = 0;
				} else {
					dp[d][j] = max (dp[d][j], max (dp[to[i]][j - 1] + a[d], dp[to[i]][j]));
				}
			}
		}
	}
	if (flag) {
		dp[d][0] = 0;
		dp[d][1] = a[d];
	}
}

main () {
	
	memset (dp, 128, sizeof (dp));
	fin >> n >> k;
	for (int i = 1; i <= n; ++ i) fin >> a[i];
	for (int i = 1; i <= n - 1; ++ i) {
		int x, l;
		fin >> x >> l;
		add (x, l);
		add (l, x);
	}
	dfs (1, 0);
	int ans = 0;
	for (int i = 0; i <= k; ++ i) ans = max (ans, dp[1][i]);
	printf ("%lld\n", ans);
	
	return 0;
}

/*
5 2
2 4 1 3 2
1 2
1 3
3 4
1 5
*/

树形dp