Submission #647482

#	Time	Username	Problem	Language	Result	Execution time	Memory
647482		ymm	Power Plant (JOI20_power)	C++17	100 / 100	225 ms	29032 KiB

power

#include <bits/stdc++.h>
#define Loop(x,l,r) for (ll x = (l); x < (ll)(r); ++x)
#define LoopR(x,l,r) for (ll x = (r)-1; x >= (ll)(l); --x)
typedef long long ll;
typedef std::pair<int, int> pii;
typedef std::pair<ll , ll > pll;
using namespace std;

const int N = 200'010;
int if_is_root[N];
bool has_gen[N];
vector<int> A[N];
int n;

int add_root_gen(int x, int rt)
{
	if (has_gen[rt])
		return max(x-1, 1);
	else
		return x;
}

void dfs1(int v, int p)
{
	int sum = 0;
	for (int u : A[v]) {
		if (u == p)
			continue;
		dfs1(u, v);
		sum += add_root_gen(if_is_root[u], u);
	}
	if_is_root[v] = sum;
}

int ans = 2;

void dfs2(int v, int p, int from_p)
{
	int sum = from_p;
	for (int u : A[v]) {
		if (u == p)
			continue;
		sum += add_root_gen(if_is_root[u], u);
	}
	ans = max(ans, add_root_gen(sum, v));
	for (int u : A[v]) {
		if (u == p)
			continue;
		int from_u = add_root_gen(if_is_root[u], u);
		dfs2(u, v, add_root_gen(sum - from_u, v));
	}
}

void handle_edge()
{
	int cnt = 0;
	Loop (i,0,n)
		cnt += has_gen[i];
	if (cnt <= 2) {
		cout << cnt << '\n';
		exit(0);
	}
}

int main()
{
	cin.tie(0) -> sync_with_stdio(false);
	cin >> n;
	Loop (i,1,n) {
		int v, u;
		cin >> v >> u;
		--v; --u;
		A[v].push_back(u);
		A[u].push_back(v);
	}
	string s;
	cin >> s;
	Loop (i,0,n)
		has_gen[i] = s[i]-'0';
	handle_edge();
	dfs1(0, -1);
	dfs2(0, -1, 0);
	cout << ans << '\n';
}

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