Submission #729735

#	Time	Username	Problem	Language	Result	Execution time	Memory
729735		pavement	Power Plant (JOI20_power)	C++17	100 / 100	201 ms	33996 KiB

power

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define int long long
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define ppb pop_back
#define eb emplace_back
#define g0(a) get<0>(a)
#define g1(a) get<1>(a)
#define g2(a) get<2>(a)
#define g3(a) get<3>(a)
#define g4(a) get<4>(a)
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
using db = double;
using ll = long long;
using ld = long double;
using ii = pair<int, int>;
using iii = tuple<int, int, int>;
using iiii = tuple<int, int, int, int>;
using iiiii = tuple<int, int, int, int, int>;
template<class key, class value = null_type, class cmp = less<key> >
using ordered_set = tree<key, value, cmp, rb_tree_tag, tree_order_statistics_node_update>;

int N, mem[200005][2];
char c;
bool S[200005];
vector<int> adj[200005];

int dp(int n, bool b, int e = -1) {
	if (mem[n][b] != -1) return mem[n][b];
	if (S[n]) {
		if (b) {
			int x = -1;
			for (auto u : adj[n]) if (u != e) {
				x += dp(u, 1, n);
			}
			return mem[n][b] = max(x, 1ll);
		} else {
			int x = 0, y = 0, z = 0;
			for (auto u : adj[n]) if (u != e) {
				x = max(x, dp(u, 1, n));
				y += dp(u, 1, n);
				z = max(z, dp(u, 0, n));
			}
			x++;
			y--;
			return mem[n][b] = max({x, y, z});
		}
	} else {
		if (b) {
			int x = 0;
			for (auto u : adj[n]) if (u != e) {
				x += dp(u, 1, n);
			}
			return mem[n][b] = max(x, 0ll);
		} else {
			int x = 0, y = 0, z = 0;
			for (auto u : adj[n]) if (u != e) {
				x = max(x, dp(u, 1, n));
				y += dp(u, 1, n);
				z = max(z, dp(u, 0, n));
			}
			return mem[n][b] = max({x, y, z});
		}
	}
}

main() {
	memset(mem, -1, sizeof mem);
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> N;
	for (int u, v, i = 1; i < N; i++) {
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for (int i = 1; i <= N; i++) {
		cin >> c;
		S[i] = (c == '1');
	}
	cout << dp(1, 0) << '\n';
}

Compilation message (stderr)

power.cpp:72:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   72 | main() {
      | ^~~~

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