Submission #1162634

#	Time	Username	Problem	Language	Result	Execution time	Memory
1162634		Lucpp	Team Coding (EGOI24_teamcoding)	C++20	100 / 100	414 ms	28940 KiB

Main

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define sz(x) (int)size(x)
#define all(x) (x).begin(), (x).end()

vector<int> color;
vector<map<int, int>> colByD;
vector<vector<int>> g;
vector<vector<int>> layerSz;
vector<int> cntAbove, cntTotal;

int maxCnt = 0, minSwaps = 0;

void dfs(int u, int par, int d){
	cntAbove[color[u]]++;
	int cntAlready = -cntTotal[color[u]];
	cntTotal[color[u]]++;
	for(int v : g[u]){
		if(v == par) continue;
		dfs(v, u, d+1);
		if(sz(layerSz[v]) > sz(layerSz[u])) swap(layerSz[u], layerSz[v]);
		for(int i = 0; i < sz(layerSz[v]); i++) layerSz[u][sz(layerSz[u])-sz(layerSz[v])+i] += layerSz[v][i];
	}
	layerSz[u].push_back(1);
	cntAbove[color[u]]--;
	cntAlready += cntTotal[color[u]];
	if(cntAbove[color[u]] == 0){
		auto it = colByD[color[u]].lower_bound(d);
		int cnt = 0;
		while(it != colByD[color[u]].end()){
			int i = sz(layerSz[u])-1 + d - it->first;
			if(i < 0) break;
			cnt += min(layerSz[u][i], it->second);
			it++;
		}
		if(cnt > maxCnt){
			maxCnt = cnt;
			minSwaps = cnt - cntAlready;
		}
		else if(cnt == maxCnt) minSwaps = min(minSwaps, cnt - cntAlready);
	}
}

void solve(){
	int n, k;
	cin >> n >> k;
	color.resize(n);
	for(int &e : color) cin >> e;
	g.resize(n);
	for(int i = 1; i < n; i++){
		int x;
		cin >> x;
		g[x].push_back(i);
	}
	colByD.resize(n);
	auto dfsDepth = [&](auto&& self, int u, int par, int d) -> void {
		colByD[color[u]][d]++;
		for(int v : g[u]){
			if(v == par) continue;
			self(self, v, u, d+1);
		}
	};
	dfsDepth(dfsDepth, 0, -1, 0);
	layerSz.resize(n);
	cntAbove.resize(k);
	cntTotal.resize(k);
	dfs(0, -1, 0);
	cout << maxCnt << " " << minSwaps << "\n";
}

int main(){
	cin.tie(0)->sync_with_stdio(false);
	solve();
}

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