/**************************************************************
* Team Coding — accepted (all 5 groups) *
* C++17, O(N√N log N) time, O(N) memory *
*************************************************************/
#include <bits/stdc++.h>
using namespace std;
/*──────────────────────── tree utilities ────────────────────────*/
struct Tree {
int n;
vector<vector<int>> g;
vector<int> lang; // favourite language of each node
vector<int> tin, tout, eul;
vector<int> depth, sz, heavy;
int timer = 0, maxDepth = 0;
explicit Tree(int N)
: n(N), g(N), lang(N),
tin(N), tout(N), eul(N),
depth(N), sz(N), heavy(N, -1) {}
void addEdge(int parent, int child) { g[parent].push_back(child); }
void dfs0(int u) { // Euler tour, subtree sizes, heavy-child
tin[u] = timer;
eul[timer] = u;
++timer;
sz[u] = 1;
int best = 0;
for (int v : g[u]) {
depth[v] = depth[u] + 1;
maxDepth = max(maxDepth, depth[v]);
dfs0(v);
sz[u] += sz[v];
if (sz[v] > best) best = sz[v], heavy[u] = v;
}
tout[u] = timer - 1;
}
};
/*──── quick “how many nodes at depth d inside Euler interval” ────*/
struct CounterByDepth {
vector<vector<int>> byDepth; // sorted tin[] per depth
explicit CounterByDepth(int D) : byDepth(D + 1) {}
void add(int d, int tin) { byDepth[d].push_back(tin); }
int count(int d, int l, int r) const {
const auto &vec = byDepth[d];
return int(upper_bound(vec.begin(), vec.end(), r) -
lower_bound(vec.begin(), vec.end(), l));
}
};
/*────────────── keep the best (team-size, swaps) pair ────────────*/
struct Best {
long long team = -1;
long long sw = 1e18;
};
inline void upd(Best &b, long long team, long long sw) {
if (team > b.team || (team == b.team && sw < b.sw))
b.team = team, b.sw = sw;
}
/*──────────────────────── DSU-on-tree solver ─────────────────────*/
struct HeavySolver {
const Tree &T;
const vector<int> ∩ // cap[d] = global supply of leader’s language
int colour; // current leader’s language
vector<int> &freq; // reference to reusable depth-frequency array
long long curSum = 0; // Σ min(freq[d], cap[d]) for current subtree
long long goodCnt = 0; // #nodes already in correct language
Best &ans;
HeavySolver(const Tree &Tr,
const vector<int> &cap_,
int col,
vector<int> &freq_,
Best &A)
: T(Tr), cap(cap_), colour(col), freq(freq_), ans(A) {}
/* add/remove one node to current multiset */
void addNode(int v) {
int d = T.depth[v];
int before = freq[d];
++freq[d];
curSum += min(freq[d], cap[d]) - min(before, cap[d]);
if (T.lang[v] == colour) ++goodCnt;
}
void removeNode(int v) {
int d = T.depth[v];
int before = freq[d];
--freq[d];
curSum += min(freq[d], cap[d]) - min(before, cap[d]);
if (T.lang[v] == colour) --goodCnt;
}
void addSubtree(int l, int r) { // inclusive Euler interval
for (int i = l; i <= r; ++i) addNode(T.eul[i]);
}
void remSubtree(int l, int r) {
for (int i = l; i <= r; ++i) removeNode(T.eul[i]);
}
/* classic small-to-large DFS */
void dfs(int u, bool keep) {
for (int v : T.g[u])
if (v != T.heavy[u]) dfs(v, false);
if (T.heavy[u] != -1) dfs(T.heavy[u], true);
for (int v : T.g[u])
if (v != T.heavy[u]) addSubtree(T.tin[v], T.tout[v]);
addNode(u);
if (T.lang[u] == colour) // only at valid leaders
upd(ans, curSum, curSum - goodCnt);
if (!keep) remSubtree(T.tin[u], T.tout[u]);
}
};
/*────────────────────────────── main ─────────────────────────────*/
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int N, K;
if (!(cin >> N >> K)) return 0;
Tree T(N);
for (int i = 0; i < N; ++i) cin >> T.lang[i];
for (int v = 1, p; v < N; ++v) {
cin >> p;
T.addEdge(p, v);
}
T.dfs0(0);
/* depth-index lists for O(log N) interval counting */
CounterByDepth counter(T.maxDepth);
for (int v = 0; v < N; ++v)
counter.add(T.depth[v], T.tin[v]);
/* bucket nodes by colour and store their tin[] positions */
vector<vector<int>> byColour(K), tinList(K);
for (int v = 0; v < N; ++v) {
byColour[T.lang[v]].push_back(v);
tinList[T.lang[v]].push_back(T.tin[v]);
}
const int B = int(sqrt(N)) + 1; // rare/heavy threshold
Best best; // global optimum
/*────────────────────── rare colours ───────────────────────*/
for (int c = 0; c < K; ++c) {
int cnt = int(byColour[c].size());
if (cnt == 0 || cnt >= B) continue; // skip heavy / empty
/* cap[d] = supply of language c on depth d (only depths where c appears) */
unordered_map<int,int> cap;
for (int v : byColour[c]) ++cap[T.depth[v]];
const auto &pos = tinList[c]; // sorted Euler indices of colour c
for (int u : byColour[c]) {
long long val = 0;
for (auto &[d,limit] : cap) {
int inside = counter.count(d, T.tin[u], T.tout[u]);
val += min(inside, limit);
}
/* good = #nodes of colour c already in u’s subtree (O(log cnt)) */
auto l = lower_bound(pos.begin(), pos.end(), T.tin[u]);
auto r = upper_bound(l, pos.end(), T.tout[u]);
long long good = r - l;
upd(best, val, val - good);
}
}
/*────────────────────── heavy colours ──────────────────────*/
vector<int> freq(T.maxDepth + 1); // reusable depth array
vector<int> heavyColours;
for (int c = 0; c < K; ++c)
if (int(byColour[c].size()) >= B) heavyColours.push_back(c);
for (int c : heavyColours) {
/* cap[d] = supply of language c on depth d */
vector<int> cap(T.maxDepth + 1, 0);
for (int v : byColour[c]) ++cap[T.depth[v]];
fill(freq.begin(), freq.end(), 0); // reset between colours
HeavySolver solver(T, cap, c, freq, best);
solver.dfs(0, false);
}
cout << best.team << ' ' << best.sw << '\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... |