Submission #952441

#	Time	Username	Problem	Language	Result	Execution time	Memory
952441		Error42	Power Plant (JOI20_power)	C++17	100 / 100	429 ms	26236 KiB

power

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

int n;
vector<vector<int>> graph;
vector<bool> removed;
vector<int> subtree_sz;
vector<bool> generator;

void calc_subtree_sz(int const cur, int const par) {
    subtree_sz[cur] = 1;

    for (int const& neigh : graph[cur]) {
        if (neigh == par || removed[neigh])
            continue;

        calc_subtree_sz(neigh, cur);
        subtree_sz[cur] += subtree_sz[neigh];
    }
}

int find_centroid(int const cur, int const par, int const tree_sz) {
    for (int const& neigh : graph[cur]) {
        if (neigh == par || removed[neigh])
            continue;

        if (subtree_sz[neigh] > tree_sz / 2)
            return find_centroid(neigh, cur, tree_sz);
    }

    return cur;
}

int calc_ans_rec(int const cur, int const par) {
    int ans = -generator[cur];

    for (int const neigh : graph[cur]) {
        if (neigh == par || removed[neigh])
            continue;

        ans += calc_ans_rec(neigh, cur);
    }

    ans = max(ans, (int)generator[cur]);

    return ans;
}

int calc_ans(int const centroid) {
    int none = generator[centroid];
    int one = 0;
    int all = -generator[centroid];

    for (int const neigh : graph[centroid]) {
        if (removed[neigh])
            continue;

        int neigh_ans = calc_ans_rec(neigh, centroid);

        one = max(one, neigh_ans + generator[centroid]);
        all += neigh_ans;
    }

    return max({ none, one, all });
}

int decomp(int const cur) {
    calc_subtree_sz(cur, -1);
    int const centroid = find_centroid(cur, -1, subtree_sz[cur]);

    int ans = calc_ans(centroid);

    removed[centroid] = true;

    for (int const neigh : graph[centroid]) {
        if (removed[neigh])
            continue;

        ans = max(ans, decomp(neigh));
    }

    return ans;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;

    graph.resize(n);
    removed.resize(n);
    subtree_sz.resize(n);
    generator.resize(n);

    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        u--; v--;

        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    string s;
    cin >> s;

    for (int i = 0; i < n; i++)
        generator[i] = s[i] == '1';

    int const ans = decomp(0);

    cout << ans << "\n";
}

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