Submission #499938

#	Time	Username	Problem	Language	Result	Execution time	Memory
499938		zhougz	Power Plant (JOI20_power)	C++17	100 / 100	155 ms	28256 KiB

power

#include <bits/stdc++.h>

using namespace std;

int n, res;
const int MAXN = 200'050;
vector<int> v[MAXN];
string s;
bool has_gen[MAXN];

int dfs(int x, int p) {
	int mx = 0, sum = 0;
	for (auto z : v[x]) {
		if (z != p) {
			int dp = dfs(z, x);
			mx = max(mx, dp);
			sum += dp;
		}
	}
	int ans;
	// My key concern was that if we assume there is
	// an ON node above x, but there actually isn't.
	// But in ans we can pretend x is the root node
	// E.g. if there is actually no ON node above x
	// then one possible case for ans would be
	// mx + 1 (ON x if has_gen[x]) (although dp is still 1)
	// Basically also calculate ans at each node assuming the root of the tree is x besides dp
	if (!has_gen[x]) {
		ans = sum;
		res = max(res, ans);
		return sum; // = dp
	}
	// Force generator ON
	ans = 1 + mx; // Assuming x is the root (i.e. no ON above x)
	int ret = 1; // Assuming has ON above x, we must pick only 1
	// Force generator OFF
	ans = max(ans, sum - 1); // Again assuming x is the root
	ret = max(ret, sum - 1); // Even if there is ON above x it's still this. Here we maximise dp before returning it
	res = max(res, ans);
	return ret;
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cin >> n;
	for (int i = 0, a, b; i < n - 1; i++) {
		cin >> a >> b;
		v[a].push_back(b);
		v[b].push_back(a);
	}
	cin >> s;
	for (int i = 1; i <= n; i++) {
		has_gen[i] = s[i - 1] == '1';
	}
	dfs(1, -1);
	cout << res;
	return 0;
}

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