제출 #1340180

#제출 시각아이디문제언어결과실행 시간메모리
1340180yogesh_saneTeam Coding (EGOI24_teamcoding)C++20
100 / 100
1029 ms32440 KiB
#include <bits/stdc++.h>
using namespace std;

// PROBLEM UNDERSTANDING:
// We have a tree where each node has a color.
// We want to find a subtree and select nodes from it (max one per depth level)
// Goal 1: Maximize total nodes selected
// Goal 2: Among optimal solutions, minimize nodes of a specific target color

const int THRESHOLD = 333; // If a color appears >= 333 times, we handle it differently

int n, k;
vector<vector<int>> children;  // children[u] = list of children of node u
vector<int> color;             // color[u] = color of node u
vector<int> depth;             // depth[u] = how deep node u is in the tree

// For checking if one node is in another's subtree
vector<int> enter_time, exit_time;  // DFS traversal times
int current_time = 0;

// For each depth level, store all enter/exit times of nodes at that depth
vector<vector<int>> enter_times_by_depth;
vector<vector<int>> exit_times_by_depth;

int best_nodes = 0;      // Maximum nodes we can select
int best_color_count = 1e9;  // Minimum target color nodes in best selection

// STEP 1: Calculate depths and DFS times
void calculate_depths(int node) {
    enter_time[node] = current_time;
    enter_times_by_depth[depth[node]].push_back(enter_time[node]);
    
    for (int child : children[node]) {
        depth[child] = depth[node] + 1;
        current_time++;
        calculate_depths(child);
    }
    
    exit_time[node] = current_time;
    exit_times_by_depth[depth[node]].push_back(exit_time[node]);
}

// Check if node 'descendant' is in the subtree of node 'ancestor'
bool is_in_subtree(int ancestor, int descendant) {
    return enter_time[descendant] >= enter_time[ancestor] && 
           exit_time[descendant] <= exit_time[ancestor];
}

// STRATEGY 1: For colors that appear MANY times (>= 333 occurrences)
// We do a special DFS to efficiently count nodes at each depth in each subtree

int target_color;
vector<int> target_count_in_subtree;  // How many target colored nodes in each node's subtree
vector<int> target_count_at_depth;    // Total target colored nodes at each depth
vector<vector<int>> nodes_at_relative_depth;  // nodes_at_relative_depth[node][d] = count at depth d below node

void count_frequent_color(int node, int root_of_interest) {
    // root_of_interest is the first ancestor with target color (or -1 if none)
    
    // If this is the first target-colored node we've seen, it becomes the root
    if (root_of_interest == -1 && color[node] == target_color) {
        root_of_interest = node;
    }
    
    // If we have a root of interest, track how far down we are from it
    if (root_of_interest != -1) {
        int relative_depth = depth[node] - depth[root_of_interest];
        
        // Expand array if needed
        if (relative_depth == nodes_at_relative_depth[root_of_interest].size()) {
            nodes_at_relative_depth[root_of_interest].push_back(0);
        }
        
        nodes_at_relative_depth[root_of_interest][relative_depth]++;
    }
    
    // Count target colored nodes
    if (color[node] == target_color) {
        target_count_in_subtree[root_of_interest]++;
        target_count_at_depth[depth[node]]++;
    }
    
    // Recurse to children
    for (int child : children[node]) {
        count_frequent_color(child, root_of_interest);
    }
}

void handle_frequent_color(vector<int>& nodes_with_color) {
    target_count_in_subtree.assign(n, 0);
    target_count_at_depth.assign(n, 0);
    nodes_at_relative_depth.assign(n, vector<int>());
    
    count_frequent_color(0, -1);
    
    // Try each node as a potential root
    for (int node = 0; node < n; node++) {
        if (target_count_in_subtree[node] == 0) continue;
        
        int total_selected = 0;
        
        // For each depth level below this node
        for (int rel_depth = 0; rel_depth < nodes_at_relative_depth[node].size(); rel_depth++) {
            int abs_depth = rel_depth + depth[node];
            int available_in_subtree = nodes_at_relative_depth[node][rel_depth];
            int total_at_this_depth = target_count_at_depth[abs_depth];
            
            // We can take minimum of: nodes in subtree OR total at this depth
            total_selected += min(available_in_subtree, total_at_this_depth);
        }
        
        // Update answer
        if (total_selected == best_nodes) {
            best_color_count = min(best_color_count, total_selected - target_count_in_subtree[node]);
        }
        if (total_selected > best_nodes) {
            best_nodes = total_selected;
            best_color_count = total_selected - target_count_in_subtree[node];
        }
    }
}

// STRATEGY 2: For colors that appear FEW times (< 333 occurrences)
// We can afford to check each node individually

void handle_rare_color(vector<int>& nodes_with_color) {
    // Count how many nodes of this color at each depth
    map<int, int> count_at_depth;
    for (int node : nodes_with_color) {
        count_at_depth[depth[node]]++;
    }
    
    // Try each colored node as the root of our selection
    for (int root : nodes_with_color) {
        int total_selected = 0;
        
        // For each depth level, count how many nodes are in root's subtree
        for (auto [d, count] : count_at_depth) {
            // Binary search: find nodes at depth d that are in root's subtree
            // A node is in subtree if: enter_time[root] <= enter_time[node] <= exit_time[root]
            
            auto start = lower_bound(enter_times_by_depth[d].begin(), 
                                    enter_times_by_depth[d].end(), 
                                    enter_time[root]);
            
            auto end = upper_bound(exit_times_by_depth[d].begin(), 
                                  exit_times_by_depth[d].end(), 
                                  exit_time[root]);
            
            int in_subtree = end - exit_times_by_depth[d].begin() - 
                           (start - enter_times_by_depth[d].begin());
            
            total_selected += min(count, in_subtree);
        }
        
        // Only compute exact count if this could be optimal
        if (total_selected >= best_nodes) {
            // Count exactly how many colored nodes in this subtree
            int colored_in_subtree = 0;
            for (int node : nodes_with_color) {
                if (is_in_subtree(root, node)) {
                    colored_in_subtree++;
                }
            }
            
            // Update answer
            if (total_selected == best_nodes) {
                best_color_count = min(best_color_count, total_selected - colored_in_subtree);
            }
            if (total_selected > best_nodes) {
                best_nodes = total_selected;
                best_color_count = total_selected - colored_in_subtree;
            }
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> n >> k;
    
    children.resize(n);
    color.resize(n);
    depth.resize(n);
    enter_time.resize(n);
    exit_time.resize(n);
    enter_times_by_depth.resize(n);
    exit_times_by_depth.resize(n);
    
    // Group nodes by their color
    vector<vector<int>> nodes_by_color(k);
    for (int i = 0; i < n; i++) {
        cin >> color[i];
        nodes_by_color[color[i]].push_back(i);
    }
    
    // Read tree structure
    for (int i = 1; i < n; i++) {
        int parent;
        cin >> parent;
        children[parent].push_back(i);
    }
    
    // Calculate depths and DFS times
    calculate_depths(0);
    
    // Process each color
    for (int c = 0; c < k; c++) {
        if (nodes_by_color[c].size() >= THRESHOLD) {
            // Many nodes of this color - use efficient method
            target_color = c;
            handle_frequent_color(nodes_by_color[c]);
        } else {
            // Few nodes of this color - use simple method
            handle_rare_color(nodes_by_color[c]);
        }
    }
    
    cout << best_nodes << ' ' << best_color_count << '\n';
    
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...