Submission #1206640

#TimeUsernameProblemLanguageResultExecution timeMemory
1206640oceanTeam Coding (EGOI24_teamcoding)C++17
31 / 100
47 ms29508 KiB
/************************************************************** * 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; // 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 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...