#include <bits/stdc++.h>
using namespace std;
//–––––– Typedefs and Macros ––––––
typedef int ll; // adjust if needed (using int for indices)
typedef pair<ll, ll> pi;
#define pb push_back
#define mp make_pair
const ll MAXN = 100005;
//–––––– Global Variables ––––––
ll n, k, label = 0;
ll color[MAXN], par[MAXN], in[MAXN], out[MAXN], cnt[MAXN], dep[MAXN], big[MAXN];
vector<ll> adj[MAXN], nodelist[MAXN];
// For each color, record all encountered depths (to later sort & unique)
vector<ll> deplist[MAXN];
// Instead of unordered_maps, we now store frequency tables as sorted vectors of (depth, frequency) pairs.
vector<vector<pi>> dp(MAXN), dp2(MAXN), dp3(MAXN), dp4(MAXN), dp5(MAXN);
// Global frequency table for each color (by depth). Index: color
vector<vector<pi>> colMapVec;
// Final answer stored as (res, -swaps) so that lexicographical max works.
pi ans = mp(-1000000000, -1000000000);
//–––––– Helper Functions for Sorted Frequency Tables –––––//
// Binary search: return frequency for key in sorted vector v, or 0 if not found.
inline int getValue(const vector<pi>& v, int key) {
int lo = 0, hi = (int)v.size() - 1;
while (lo <= hi) {
int mid = (lo + hi) >> 1;
if (v[mid].first == key)
return v[mid].second;
else if (v[mid].first < key)
lo = mid + 1;
else
hi = mid - 1;
}
return 0;
}
// Insert or update key in sorted vector v by adding val.
inline void addValue(vector<pi>& v, int key, int val) {
auto it = std::lower_bound(v.begin(), v.end(), mp(key, 0));
if (it != v.end() && it->first == key)
it->second += val;
else
v.insert(it, mp(key, val));
}
// Merge two sorted frequency vectors: merge src into dest.
inline void mergeVectors(vector<pi>& dest, const vector<pi>& src) {
vector<pi> merged;
merged.reserve(dest.size() + src.size());
int i = 0, j = 0;
while (i < (int)dest.size() && j < (int)src.size()){
if (dest[i].first == src[j].first) {
merged.pb(mp(dest[i].first, dest[i].second + src[j].second));
i++; j++;
} else if (dest[i].first < src[j].first) {
merged.pb(dest[i]);
i++;
} else {
merged.pb(src[j]);
j++;
}
}
while (i < (int)dest.size()){
merged.pb(dest[i]); i++;
}
while (j < (int)src.size()){
merged.pb(src[j]); j++;
}
dest = move(merged);
}
//–––––– Iterative DFS for Euler Tour (Computes in[], out[], dep[]) –––––//
void iterativeDFS_precomp1() {
vector<int> st;
vector<int> ptr(n, 0); // pointer for each node's child index
dep[0] = 0;
st.pb(0);
while (!st.empty()) {
int u = st.back();
// First time visiting u: record entry time and add its depth to deplist.
if (ptr[u] == 0) {
in[u] = ++label;
deplist[color[u]].pb(dep[u]);
}
if (ptr[u] < (int)adj[u].size()) {
int v = adj[u][ptr[u]++];
dep[v] = dep[u] + 1;
st.pb(v);
} else {
out[u] = label;
st.pop_back();
}
}
}
//–––––– Get Postorder Order (Iterative) –––––//
vector<int> getPostorder() {
vector<int> post;
vector<int> st;
st.pb(0);
while(!st.empty()){
int u = st.back();
st.pop_back();
post.pb(u);
for (int v : adj[u])
st.pb(v);
}
reverse(post.begin(), post.end());
return post;
}
//–––––– Precomp2: DSU-on-Tree Merging Using Sorted Vectors (Iterative Postorder) –––––//
void precomp2_iterative() {
// Initialize global frequency table for each color.
colMapVec.assign(k, vector<pi>());
for (int u = 0; u < n; u++) {
addValue(colMapVec[color[u]], dep[u], 1);
}
vector<int> post = getPostorder();
// Process nodes in postorder (children before parent).
for (int u : post) {
// Merge each child's dp into dp[u].
for (int v : adj[u]) {
if (dp[u].size() < dp[v].size())
swap(dp[u], dp[v]);
mergeVectors(dp[u], dp[v]);
dp[v].clear();
}
// Add the current node.
addValue(dp[u], dep[u], 1);
// For “small” colors, build dp2[u] using only the depths that appear for that color.
if (!big[color[u]]) {
for (int d : deplist[color[u]]) {
int freq = getValue(dp[u], d);
if (freq > 0)
addValue(dp2[u], d, freq);
}
}
}
}
//–––––– Solve for “big” colors using Iterative Postorder Merging for dp4/dp5 –––––//
void solve_big_iterative(int c) {
// Build a "yes" array: yes[u] is true if none of u's ancestors (excluding u) has color c.
vector<bool> yes(n, false);
{
vector<int> st;
vector<int> ptr(n, 0);
yes[0] = true;
st.pb(0);
while(!st.empty()){
int u = st.back();
if(ptr[u] < (int)adj[u].size()){
int v = adj[u][ptr[u]++];
yes[v] = yes[u] && (color[u] != c);
st.pb(v);
} else {
st.pop_back();
}
}
}
vector<int> post = getPostorder();
// Clear dp4 and dp5 for all nodes.
for (int u = 0; u < n; u++){
dp4[u].clear();
dp5[u].clear();
}
// Process nodes in postorder.
for (int u : post) {
if (yes[u] && color[u] == c) {
int res = 0, res2 = 0;
// For every unique depth for color c deeper than dep[u], compute contributions.
for (int d : deplist[c]) {
if(d <= dep[u]) continue;
int a = getValue(dp4[u], d);
int b = getValue(colMapVec[c], d);
int take = (a < b ? a : b);
res += take;
int cval = getValue(dp5[u], d);
res2 += (take > cval ? take - cval : 0);
}
ans = max(ans, mp(res, -res2));
}
// Add the current node u.
addValue(dp4[u], dep[u], 1);
if(c == color[u])
addValue(dp5[u], dep[u], 1);
// Merge u's dp4/dp5 into its parent (if not root).
if(u != 0) {
int p = par[u];
if (dp4[p].size() < dp4[u].size())
swap(dp4[p], dp4[u]);
mergeVectors(dp4[p], dp4[u]);
dp4[u].clear();
if (dp5[p].size() < dp5[u].size())
swap(dp5[p], dp5[u]);
mergeVectors(dp5[p], dp5[u]);
dp5[u].clear();
}
}
}
//–––––– Solve for “small” colors –––––//
void solve_small(int c) {
int sz = (int)nodelist[c].size();
// For every pair of nodes of color c (where one is an ancestor of the other),
// update dp3 for the ancestor.
for (int i = 0; i < sz; i++) {
for (int j = i + 1; j < sz; j++) {
int x = nodelist[c][i], y = nodelist[c][j];
if (in[x] > in[y])
swap(x, y);
if (in[x] <= in[y] && in[y] <= out[x])
addValue(dp3[x], dep[y], 1);
}
}
// For each node of color c, combine dp2 and dp3 to update the candidate answer.
for (int i = 0; i < sz; i++) {
int x = nodelist[c][i];
int res = 0, res2 = 0;
for (int d : deplist[c]) {
if(d <= dep[x]) continue;
int a = getValue(dp2[x], d);
int b = getValue(colMapVec[c], d);
int take = (a < b ? a : b);
res += take;
int cval = getValue(dp3[x], d);
res2 += (take > cval ? take - cval : 0);
}
dp2[x].clear();
dp3[x].clear();
ans = max(ans, mp(res, -res2));
}
}
//–––––– Main Function –––––//
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
cout.tie(nullptr);
cin >> n >> k;
// Resize global frequency table vector for each color.
colMapVec.resize(k);
// Read node colors and fill nodelist and frequency counts.
for (int i = 0; i < n; i++){
cin >> color[i];
cnt[color[i]]++;
nodelist[color[i]].pb(i);
}
// Mark "big" colors: here using a threshold (blkThreshold).
const int blkThreshold = 316; // can adjust threshold as needed
for (int i = 0; i < k; i++){
if (cnt[i] > blkThreshold)
big[i] = 1;
}
// Read parent pointers and build tree. (Assuming input is 0-indexed.)
for (int i = 1; i < n; i++){
cin >> par[i];
adj[par[i]].pb(i);
}
// Compute Euler Tour iteratively: fill in[], out[], dep[] and deplist.
iterativeDFS_precomp1();
// For each color, sort and unique the deplist.
for (int i = 0; i < k; i++){
sort(deplist[i].begin(), deplist[i].end());
deplist[i].erase(unique(deplist[i].begin(), deplist[i].end()), deplist[i].end());
}
// Build dp and dp2 using DSU-on-tree merging (iterative postorder).
precomp2_iterative();
// Process each color: if color is "big" use iterative DFS merging (dp4/dp5), else use solve_small.
for (int c = 0; c < k; c++){
if (big[c]) {
solve_big_iterative(c);
} else {
solve_small(c);
}
}
// Output final answer (note: answer is stored as (res, -swaps) so we output ans.first+1 and -ans.second).
cout << ans.first + 1 << " " << -ans.second << "\n";
return 0;
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |