Submission #1098666

#	Time	Username	Problem	Language	Result	Execution time	Memory
1098666		Trisanu_Das	Team Coding (EGOI24_teamcoding)	C++17	62 / 100	4096 ms	30668 KiB

Main

#include <bits/stdc++.h>

using namespace std;

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

    int n, k;
    cin >> n >> k;

    vector<int> c(n);
    for (int i = 0; i < n; ++i) {
        cin >> c[i];
    }

    vector<vector<int>> g(n);
    for (int x = 1; x < n; ++x) {
        int p;
        cin >> p;
        g[p].emplace_back(x);
    }

    vector<int> level(n), is_root(n), below_same_color(n, 1);
    {
        vector<int> col_parent(n, -1);
        function<void(int)> dfs = [&](int x) {
            int cp = col_parent[c[x]];
            col_parent[c[x]] = x;
            for (int y : g[x]) {
                level[y] = level[x] + 1;
                dfs(y);
            }
            col_parent[c[x]] = cp;
            if (cp == -1) {
                is_root[x] = true;
            } else {
                below_same_color[cp] += below_same_color[x];
            }
        };
        dfs(0);
    }

    pair<int, int> best = {1, 1};

    vector<map<int, int>> col_at_lev(n);
    for (int i = 0; i < n; ++i) {
        ++col_at_lev[c[i]][level[i]];
    }

    {
        vector<vector<int>> below(n);
        function<void(int)> dfs = [&](int x) {
            for (int y : g[x]) {
                dfs(y);
                if (below[y].size() > below[x].size()) {
                    swap(below[x], below[y]);
                }
                for (int i = 0; i < below[y].size(); ++i) {
                    below[x][below[x].size() - 1 - i] += below[y][below[y].size() - 1 - i];
                }
            }

            below[x].emplace_back(1);
            if (!is_root[x]) return;

            int ans = 1;
            for (auto it = col_at_lev[c[x]].upper_bound(level[x]); it != col_at_lev[c[x]].end(); ++it) {
                auto [lev, cnt] = *it;
                int idx = below[x].size() - (lev - level[x]) - 1;
                if (idx < 0) continue;
                ans += min(below[x][idx], cnt);
            }
            best = max(best, {ans, below_same_color[x]});

        };
        dfs(0);
    }

    cout << best.first << " " << best.first - best.second << endl;
}

Compilation message (stderr)

Main.cpp: In lambda function:
Main.cpp:60:35: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   60 |                 for (int i = 0; i < below[y].size(); ++i) {
      |                                 ~~^~~~~~~~~~~~~~~~~

• Subtask 1: 12 / 12
• Subtask 2: 0 / 19
• Subtask 3: 27 / 27
• Subtask 4: 23 / 23
• Subtask 5: 0 / 19