Submission #887609

#	Time	Username	Problem	Language	Result	Execution time	Memory
887609		box	Power Plant (JOI20_power)	C++17	100 / 100	131 ms	39520 KiB

power

/*
dp[i][0/1][0/1] = maximum answer for subtree of i with fixed (outside of subtree) & (subtree) existence of an activated generator
*/
#include <bits/stdc++.h>
using namespace std;

#define ar array

const int N = 2e5;

int n;
vector<int> adj[N];
string s;
int dp[N][2][2];

ar<int, 2> one(const ar<int, 2> one, const ar<int, 2> two) {
    return {one[0] + two[0], max(one[1] + two[1], max(one[0] + two[1], one[1] + two[0]))};
}
ar<int, 2> two(const ar<int, 2> one, const ar<int, 2> two) {
    return {one[0] + two[0], max(one[1] + two[0], one[0] + two[1])};
}
ar<int, 3> three(const ar<int, 3> one, const ar<int, 2> two) {
    ar<int, 3> me{-n, -n, -n};
    for (int a = 0; a < 3; a++) for (int b = 0; b < 2; b++) me[min(2, a + b)] = max(me[min(2, a + b)], one[a] + two[b]);
    return me;
}

void dfs(int p, int i) {
    for (int j : adj[i]) if (p != j) dfs(i, j);
    ar<int, 3> f{0, -n, -n};
    ar<int, 2> f2{0, -n};
    int z = 0;
    for (int j : adj[i]) if (p != j) {
        z += dp[j][0][0];
        f = three(f, ar<int, 2>{dp[j][1][0], dp[j][1][1]});
        f2 = two(f2, ar<int, 2>{dp[j][1][0], dp[j][0][1]});
    }
    dp[i][1][0] = f[0];
    dp[i][1][1] = max(f[1] - (s[i] == '1'), f[2] - (s[i] == '1'));
    dp[i][0][0] = z;
    dp[i][0][1] = max(f2[1], f[2] - (s[i] == '1'));
    if (s[i] == '1') {
        dp[i][0][1] = max(dp[i][0][1], max({f[0] + 1, f[1] + 1, f[2] - 1}));
        dp[i][1][1] = max(dp[i][1][1], max({f[0] + 1, f[1] - 1, f[2] - 1}));
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n;
    for (int h = 0; h < n - 1; h++) {
        int i, j;
        cin >> i >> j, i--, j--;
        adj[i].push_back(j);
        adj[j].push_back(i);
    }
    cin >> s;
    dfs(-1, 0);
    cout << max(dp[0][0][0], dp[0][0][1]) << '\n';
}

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