DP题解 | #小美的树上染色#

提供一种DP的思路；

dp[i][0/1]表示第i个结点染色不染色/染色的最大的染色数量。

那么对于dp[i][0]来说就是它的子节点的染色最大树量的和，dp[i][1]则需要再来一轮循环，统计某个子节点与父节点一起染色的最大值。

#include <bits/stdc++.h>

using namespace std;

using i64 = long long;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);

	int n;
	cin >> n;

	vector a(n + 1, 0);
	for (int i = 1; i <= n; i ++) {
		cin >> a[i];
	}

	vector g(n + 1, vector<int>());
	for (int i = 1; i < n; i ++) {
		int u, v;
		cin >> u >> v;
		g[u].emplace_back(v);
		g[v].emplace_back(u);
	}

	vector dp(n + 1, array<int, 2>());

	auto dfs = [&]( auto && self, int u, int fa) -> void{
		for (const auto &v : g[u]) {
			if (v == fa) continue;
			self(self, v, u);
			dp[u][0] += max(dp[v][0], dp[v][1]);
		}
		for (const auto &v : g[u]) {
			if (v == fa)continue;
			i64 x = sqrt(1LL * a[u] * a[v]);
			if (x * x == 1LL * a[u] * a[v]) {
				dp[u][1] = max(dp[u][1], dp[u][0] - max(dp[v][0], dp[v][1]) + dp[v][0] + 2);
			}
		}
	};

	dfs(dfs, 1, -1);

	cout << max(dp[1][0], dp[1][1]) << "\n";

	return 0;
}