Submission #1004732

#	Time	Username	Problem	Language	Result	Execution time	Memory
1004732		julian	Power Plant (JOI20_power)	C++98	0 / 100	1 ms	348 KiB

power

#include <bits/stdc++.h>

long dfs(long v, long parent, std::vector<std::vector<long>>& adj, std::vector<bool>& b, std::vector<std::pair<long, long>>& dp, bool blocked) {
	if (!blocked && dp[v].first != std::numeric_limits<long>::max()) return dp[v].first;
	if (blocked && dp[v].second != std::numeric_limits<long>::max()) return dp[v].second;

	dp[v].first = 0;
	dp[v].second = 0;

	if (b[v - 1]) {
		long max = 0;

		for (long e : adj[v]) {
			if (e == parent) continue;
			dp[v].first += dfs(e, v, adj, b, dp, 0);
			dp[v].second = std::max(dp[v].second, dfs(e, v, adj, b, dp, 1));
			max = std::max(max, dp[v].first);
		}

		dp[v].first = std::max(dp[v].first - 1l, 1l);
		dp[v].second = std::max(dp[v].second - 1l, max + 1l);

	} else {
		// Power plant doesn't matter
		dp[v].first = 0;
		dp[v].second = 0;
		for (long e : adj[v]) {
			if (e == parent) continue;
			dp[v].first += dfs(e, v, adj, b, dp, 0);
			dp[v].second = std::max(dp[v].second, dfs(e, v, adj, b, dp, 1));
		}
	}

	return blocked ? dp[v].second : dp[v].first;
}

int main() {
	long n;
	std::cin >> n;

	std::vector<std::vector<long>> adj(n + 1, std::vector<long>());

	std::vector<std::pair<long, long>> dp(n + 1, std::pair<long, long>{std::numeric_limits<long>::max(), std::numeric_limits<long>::max()});

	for (long i = 0; i < n - 1; i++) {
		long a, b;
		std::cin >> a >> b;
		adj[a].push_back(b);
		adj[b].push_back(a);
	}

	std::string s;
	std::cin >> s;

	std::vector<bool> b(n);
	for (size_t i = 0; i < n; i++) {
		b[i] = s[i] == '1';
	}

	dfs(1, 0, adj, b, dp, 0);
	dfs(1, 0, adj, b, dp, 1);

	long m = 0;

	for (size_t i = 1; i < dp.size(); i++) {
		m = std::max({m, dp[i].first, dp[i].second});
	}

	std::cout << m << std::endl;
}

Compilation message (stderr)

power.cpp: In function 'int main()':
power.cpp:56:23: warning: comparison of integer expressions of different signedness: 'size_t' {aka 'long unsigned int'} and 'long int' [-Wsign-compare]
   56 |  for (size_t i = 0; i < n; i++) {
      |                     ~~^~~

• Subtask 1: 0 / 6
• Subtask 2: 0 / 41
• Subtask 3: 0 / 53