Submission #1188375

#	Time	Username	Problem	Language	Result	Execution time	Memory
1188375		esomer	Team Coding (EGOI24_teamcoding)	C++20	100 / 100	163 ms	132592 KiB

Main

#include <bits/stdc++.h>

using namespace std;

const int B = 300;

vector<int> bigColors; //colors with more than B people.
vector<int> anc; //the highest ancestor of each color.
vector<vector<int>> all; //all the people with each color.
vector<int> depth, maxDepth;
vector<int> cntDepth, cntColor;
vector<vector<int>> depthCol; //the number of nodes at each depth for the bigColors.
pair<int, int> ans = {0, 0};

void getDepth(int x, vector<vector<int>>& adj) {
	maxDepth[x] = depth[x];
	for (int node : adj[x]) {
		depth[node] = depth[x] + 1;
		getDepth(node, adj);
		maxDepth[x] = max(maxDepth[x], maxDepth[node]);
	}
}

void DFS(int x, vector<vector<int>>& adj, vector<int>& l) {
	int c = l[x];
	cntDepth[depth[x]]++;
	cntColor[c]++;
	if ((int)all[c].size() <= B) {
		if (anc[c] == -1) {
			anc[c] = x;
			//I iterate over all of them.
			int cntColorBefore = cntColor[c];
			vector<int> before((int)all[c].size(), 0); //the number of nodes of each interesting depth before, i don't repeat depths.
			for (int i = 0; i < (int)all[c].size(); i++) {
				int y = all[c][i];
				if (depth[y] <= depth[x] || (i > 0 && depth[y] == depth[all[c][i-1]])) continue;
				before[i] = cntDepth[depth[y]];
			}
			for (int node : adj[x]) {
				DFS(node, adj, l);
			}
			pair<int, int> res = {1, 0};
			for (int i = 0; i < (int)all[c].size(); i++) {
				int y = all[c][i];
				if (depth[y] <= depth[x] || (i > 0 && depth[y] == depth[all[c][i-1]])) continue;
				int total = cntDepth[depth[y]] - before[i]; //total #nodes with depth[y] in the subtree.
				int numColor = 1; //total #node of color c with depth[y].
				int mx = i;
				for (int j = i+1; j < (int)all[c].size(); j++) {
					if (depth[all[c][j]] == depth[y]) numColor++;
					else break;
					mx = j;
				}
				i = mx;
				res.first += min(numColor, total);
			}
			res.second = res.first - 1 - (cntColor[c] - cntColorBefore);
			if (res.first > ans.first) ans = res;
			else if (res.first == ans.first) ans.second = min(ans.second, res.second);
			anc[c] = -1;
		} else {
			for (int node : adj[x]) {
				DFS(node, adj, l);
			}
		}
	} else {
		if (anc[c] == -1) {
			anc[c] = x;
			vector<int> before;
			for (int i = depth[x]+1; i <= maxDepth[x]; i++) {
				before.push_back(cntDepth[i]);
			}
			int cntColorBefore = cntColor[c];
			for (int node : adj[x]) {
				DFS(node, adj, l);
			}
			pair<int, int> res = {1, 0};
			for (int i = depth[x] + 1; i <= maxDepth[x]; i++) {
				int total = cntDepth[i] - before[i - (depth[x] + 1)];
				int numColor = depthCol[c][i];
				res.first += min(total, numColor);
			}
			res.second = res.first - 1 - (cntColor[c] - cntColorBefore);
			if (res.first > ans.first) ans = res;
			else if (res.first == ans.first) ans.second = min(ans.second, res.second);
			anc[c] = -1;
		} else {
			for (int node : adj[x]) {
				DFS(node, adj, l);
			}
		}
	}
	//I still have to update anc.
}

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

	int N, K; cin >> N >> K;
	vector<int> l(N);
	all.resize(K);
	for (int i = 0; i < N; i++) {
		cin >> l[i];
		all[l[i]].push_back(i);
	}
	for (int i = 0; i < K; i++) {
		if ((int)all[i].size() > B) bigColors.push_back(i);
	}
	vector<vector<int>> adj(N);
	for (int i = 1; i < N; i++) {
		int p; cin >> p;
		adj[p].push_back(i);
	}
	depth.assign(N, 0);
	maxDepth.assign(N, 0);
	getDepth(0, adj);
	depthCol.resize(K);
	for (int i = 0; i < K; i++) {
		vector<pair<int, int>> nw;
		for (int x : all[i]) nw.push_back({depth[x], x});
		sort(nw.begin(), nw.end());
		all[i].clear();
		for (pair<int, int> pr : nw) all[i].push_back(pr.second);
		if ((int)all[i].size() > B) {
			depthCol[i].assign(N, 0);
			for (int x : all[i]) depthCol[i][depth[x]]++;
		}
	}
	anc.assign(K, -1);
	cntDepth.assign(N, 0);
	cntColor.assign(K, 0);
	DFS(0, adj, l);
	cout << ans.first << " " << ans.second << "\n";
}

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