Submission #1274225

#	Time	Username	Problem	Language	Result	Execution time	Memory
1274225		trufanov.p	Team Coding (EGOI24_teamcoding)	C++20	62 / 100	4094 ms	47552 KiB

Main

#include <iostream>
#include <vector>
#include <algorithm>
#include <set>
#include <random>
#include <map>
#include <unordered_map>
#include <cmath>

using namespace std;

typedef long long ll;

#pragma GCC optimize("O3")
#pragma GCC optimization("Ofast,unroll-loops")
#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")

const int SIZE = 1e5 + 5;

int n, k;
int C;
int c[SIZE];
int cnt[SIZE];
vector<int> light, heavy;
vector<int> light_sets[SIZE];
vector<int> gr[SIZE];

int tin[SIZE], tout[SIZE];
int timer = 0;
vector<int> atDepth[SIZE];
int d[SIZE];

unordered_map<int, vector<int>> colOnDepth[SIZE];

int maxSubTree = 0, minSwaps = 0;

void dfs_precalc(int v, int cur_d = 0) {
    tin[v] = timer++;
    atDepth[cur_d].push_back(tin[v]);
    colOnDepth[c[v]][cur_d].push_back(tin[v]);
    d[v] = cur_d;
    for (int u : gr[v]) {
        dfs_precalc(u, cur_d + 1);
    }
    tout[v] = timer++;
}

int countSubDepth(int v, int d) {
    return lower_bound(atDepth[d].begin(), atDepth[d].end(), tout[v]) - lower_bound(atDepth[d].begin(), atDepth[d].end(), tin[v]);
}

int countSubColDepth(int v, int d) {
    return lower_bound(colOnDepth[c[v]][d].begin(), colOnDepth[c[v]][d].end(), tout[v]) - lower_bound(colOnDepth[c[v]][d].begin(), colOnDepth[c[v]][d].end(), tin[v]);
}

void solve_light(int c) {
    for (int v : light_sets[c]) {
        int cur_ans = 1, cur_swaps = 0;
        for (const auto& p : colOnDepth[c]) {
            if (p.first <= d[v]) {
                continue;
            }
            int poss = min(countSubDepth(v, p.first), (int)p.second.size());
            cur_ans += poss;
            cur_swaps += poss - countSubColDepth(v, p.first);
        }
        if (cur_ans > maxSubTree || (cur_ans == maxSubTree && cur_swaps < minSwaps)) {
            maxSubTree = cur_ans;
            minSwaps = cur_swaps;
        }
    }
}

void solve_heavy(int v) {
    int cur_ans = 1, cur_swaps = 0;
    for (int dep = d[v] + 1; dep < n; ++dep) {
        int poss = min(countSubDepth(v, dep), (int)colOnDepth[c[v]][dep].size());
        cur_ans += poss;
        cur_swaps += poss - countSubColDepth(v, dep);
    }
    if (cur_ans > maxSubTree || (cur_ans == maxSubTree && cur_swaps < minSwaps)) {
        maxSubTree = cur_ans;
        minSwaps = cur_swaps;
    }
}

void dfs_heavy(int v, int col) {
    if (c[v] == col) {
        solve_heavy(v);
        return;
    }
    for (int u : gr[v]) {
        dfs_heavy(u, col);
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> k;
    C = (int)sqrt(n);
    for (int i = 0; i < n; ++i) {
        cin >> c[i];
        cnt[c[i]]++;
    }
    for (int i = 0; i < k; ++i) {
        if (cnt[i] != 0) {
            if (cnt[i] < C) {
                light.push_back(i);
            } else {
                heavy.push_back(i);
            }
        }
    }
    for (int i = 0; i < n; ++i) {
        if (cnt[c[i]] < C) {
            light_sets[c[i]].push_back(i);
        }
    }
    for (int i = 1; i < n; ++i) {
        int p;
        cin >> p;
        gr[p].push_back(i);
    }
    dfs_precalc(0);
    for (int c : light) {
        solve_light(c);
    }
    for (int c : heavy) {
        dfs_heavy(0, c);
    }
    cout << maxSubTree << ' ' << minSwaps << '\n';
}

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