| # | Time | Username | Problem | Language | Result | Execution time | Memory |
|---|---|---|---|---|---|---|---|
| 1206634 | ocean | Team Coding (EGOI24_teamcoding) | C++17 | 0 ms | 0 KiB |
#include <iostream>
#include <vector>
#include <map>
#include <algorithm> // For std::min, std::max, std::swap
using namespace std;
// Global variables for tree structure and properties
vector<vector<int>> adj; // Adjacency list for children (representing the hierarchy)
vector<int> languages; // Preferred language for each employee (indexed by employee ID)
vector<int> depth_arr; // Depth of each employee from Anneke (CEO 0 at depth 0)
// Global variables to store the best results found
int max_P = 0; // Maximum number of employees in the project team
int min_S = 0; // Minimum number of switches needed to achieve max_P
// Global variable to store total counts of (language, depth) pairs for the entire company
// total_depth_lang_counts[d][k] stores count of employees at depth 'd' with language 'k'
// Using map<int, int> for the inner structure handles sparse language IDs efficiently.
vector<map<int, int>> total_depth_lang_counts;
// DFS to compute the depth of each employee in the tree
// u: current employee ID
// d: current depth
void compute_depths_dfs(int u, int d) {
depth_arr[u] = d; // Set depth for current employee
for (int v : adj[u]) { // Iterate through subordinates (children)
compute_depths_dfs(v, d + 1); // Recurse for children at next depth
}
}
// DFS function that uses small-to-large merging to compute subtree statistics
// It returns a pointer to a map representing (depth -> (language -> count)) for the subtree rooted at 'u'.
// The caller takes ownership of the returned pointer and is responsible for its deletion or merging.
map<int, map<int, int>>* dfs_solve(int u) {
// Create a new map for this node's direct contribution
map<int, map<int, int>>* u_subtree_counts = new map<int, map<int, int>>();
// Add the current node's language at its depth to its subtree counts
(*u_subtree_counts)[depth_arr[u]][languages[u]]++;
// Process children
for (int v : adj[u]) {
// Recursively get subtree counts for child 'v'
map<int, map<int, int>>* v_subtree_counts = dfs_solve(v);
// Small-to-large merging logic: always merge the smaller map into the larger one.
// This optimizes the overall time complexity.
if (u_subtree_counts->size() < v_subtree_counts->size()) {
swap(u_subtree_counts, v_subtree_counts); // 'u_subtree_counts' now points to the larger map
}
// Merge 'v_subtree_counts' into 'u_subtree_counts'
for (auto const& [d_val, lang_map_at_d] : *v_subtree_counts) { // Iterate through depths in smaller map
for (auto const& [lang_val, count] : lang_map_at_d) { // Iterate through languages at each depth
(*u_subtree_counts)[d_val][lang_val] += count; // Add count to the larger map
}
}
delete v_subtree_counts; // Free memory of the merged (smaller) map
}
// Now, 'u_subtree_counts' contains the complete (depth -> (language -> count))
// statistics for the subtree rooted at 'u'.
// Calculate P and S if 'u' is chosen as the team lead.
int current_team_lead_lang = languages[u]; // Project language is team lead's preferred language [cite: 10]
int p_for_u = 0; // Initial team members within u's subtree with target language
int s_for_u = 0; // Switches needed for u as team lead
// Iterate through all depths present in u's subtree
for (auto const& [d, lang_map_at_d] : *u_subtree_counts) {
// P_initial: Count of employees in u's subtree preferring the target_lang at depth 'd'
p_for_u += lang_map_at_d[current_team_lead_lang];
// Calculate 'X_d': employees in u's subtree at depth 'd' NOT preferring target_lang
int in_subtree_wrong_lang_at_d = 0;
for (auto const& [lang, count] : lang_map_at_d) {
if (lang != current_team_lead_lang) {
in_subtree_wrong_lang_at_d += count; // These can be swapped out [cite: 12]
}
}
// Calculate 'Y_d': employees NOT in u's subtree at depth 'd' who DO prefer target_lang [cite: 13]
int total_same_lang_at_d_globally = 0;
// Check if the language exists at this depth globally before accessing
if (total_depth_lang_counts[d].count(current_team_lead_lang)) {
total_same_lang_at_d_globally = total_depth_lang_counts[d][current_team_lead_lang];
}
int in_subtree_same_lang_at_d = lang_map_at_d[current_team_lead_lang]; // Employees in u's subtree with target language
int out_of_subtree_correct_lang_at_d = total_same_lang_at_d_globally - in_subtree_same_lang_at_d;
// Add to switches: min(X_d, Y_d) possible swaps at this depth [cite: 12, 13, 14]
s_for_u += min(in_subtree_wrong_lang_at_d, out_of_subtree_correct_lang_at_d);
}
int current_P = p_for_u + s_for_u; // Total team size for u as team lead
int current_S = s_for_u; // Total switches for u as team lead
// Update global maximum P and minimum S
if (current_P > max_P) {
max_P = current_P;
min_S = current_S;
} else if (current_P == max_P) { // If same P, choose the one with fewer switches [cite: 25]
min_S = min(min_S, current_S);
}
return u_subtree_counts; // Return the calculated subtree counts for 'u'
}
int main() {
// Optimize C++ standard streams for faster input/output
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int N, K; // N: number of employees, K: number of programming languages [cite: 21]
cin >> N >> K;
languages.resize(N);
for (int i = 0; i < N; ++i) {
cin >> languages[i]; // Read preferred language for each employee [cite: 22]
}
// Build adjacency list. Employee 'i' reports to 'b_i', so 'b_i' is parent of 'i'.
adj.resize(N);
for (int i = 1; i < N; ++i) { // Employees 1 to N-1 have bosses [cite: 24]
int b;
cin >> b; // Direct boss of employee 'i' [cite: 23]
adj[b].push_back(i); // Add 'i' as a child of 'b'
}
// Compute depths for all employees, starting with Anneke (CEO 0) at depth 0
depth_arr.resize(N);
compute_depths_dfs(0, 0);
// Precompute total counts of (language, depth) for the entire company.
// Max possible depth is N-1.
total_depth_lang_counts.resize(N);
for (int i = 0; i < N; ++i) {
total_depth_lang_counts[depth_arr[i]][languages[i]]++;
}
// Start the main DFS to find the maximum P and minimum S.
// The pointer returned for the root (employee 0) needs to be deleted to clean up memory.
map<int, map<int, int>>* final_subtree_counts_root = dfs_solve(0);
delete final_subtree_counts_root; // Clean up memory
// Output the results [cite: 25]
cout << max_P << " " << min_S << endl;
return 0;
}