#include <bits/stdc++.h>
#define fi first
#define se second
#define pb push_back
#define eb emplace_back
using namespace std;
typedef long long LL;
typedef pair<int,int> pii;
typedef vector<int> vi;
template<class T> inline T re(){
T N = 0; char c = getchar(); bool neg = 0;
for (; c < '0' || c > '9'; c = getchar()) neg |= c == '-';
for (; c >= '0' && c <= '9'; c = getchar())
N = (N<<3)+(N<<1) + c - '0';
return neg ? -N : N;
}
const int SQRT = 325;
const int MX = 1e5;
int n, k;
int c[MX + 5], p[MX + 5], col_count[MX + 5];
int dep[MX + 5], at_dep[MX + 5], max_dep = 0;
vi chld[MX + 5];
int cur_col[MX + 5]; // at each dep, # with current color
int cur_dep[MX + 5], cur_at[MX + 5];
int ans_tot = 0, ans_moves = 0;
void process_whole_tree(int col) {
for (int i = 0; i < n; i++) {
if (c[i] == col) cur_col[dep[i]]++;
}
queue<int> q;
q.push(0);
while (!q.empty()) {
int u = q.front(); q.pop();
cerr << "cur : " << u << " col: " << col << endl;
if (c[u] != col) {
for (int nx : chld[u]) {
q.push(nx);
}
} else {
// process subtree rooted at u
queue<int> tmp; tmp.push(u);
vi in_subtree;
while (!tmp.empty()) {
int nx = tmp.front(); tmp.pop();
cur_dep[dep[nx]]++;
if (c[nx] == col) {
cur_at[dep[nx]]++;
}
for (int nxnx : chld[nx]) tmp.push(nxnx);
in_subtree.pb(nx);
}
cerr << "in_subtree: "; for (int i : in_subtree) cerr << "(" << i << ": " << dep[i] << ") "; cerr << '\n';
cerr << "cur_dep: "; for (int i = dep[u]; i <= n; i++) cerr << "(" << i << ": " << cur_dep[i] << ") "; cerr << '\n';
cerr << "cur_at: "; for (int i = dep[u]; i <= n; i++) cerr << "(" << i << ": " << cur_at[i] << ") "; cerr << '\n';
cerr << "cur_col: "; for (int i = dep[u]; i <= n; i++) cerr << "(" << i << ": " << cur_col[i] << ") "; cerr << '\n';
int moves = 0, cur_tot = 1; // start with 1 as root should be counted.
for (int i = dep[u] + 1; i <= max_dep; i++) {
int tmptmp = min(cur_dep[i], cur_col[i]);
cur_tot += tmptmp;
moves += tmptmp - cur_at[i];
}
cerr << "subtree " << u << " color " << col << ": " << cur_tot << ' ' << moves << '\n' << endl;
if (cur_tot > ans_tot) {
ans_tot = cur_tot, ans_moves = moves;
} else if (cur_tot == ans_tot) {
ans_moves = min(ans_moves, moves);
}
cur_tot = 0, moves = 0;
tmp.push(u);
while (!tmp.empty()) {
int nx = tmp.front(); tmp.pop();
cur_dep[dep[nx]]--;
if (c[nx] == col) {
cur_at[dep[nx]]--;
}
for (int nxnx : chld[nx]) tmp.push(nxnx);
}
}
}
for (int i = 0; i < n; i++) {
if (c[i] == col) cur_col[dep[i]]--;
}
}
int main() {
/**
* for each node x, calculate:
* 1) number of nodes y with dep[y] > x with caveat each level has max card(z in subtree x with dep[z] fixed)
* 2) min number of swaps
*/
n = re<int>(); k = re<int>();
for (int i = 0; i < n; i++) col_count[c[i] = re<int>()]++;
for (int i = 1; i < n; i++) {
p[i] = re<int>();
chld[p[i]].pb(i);
}
[&]() {
stack<int> stk; stk.push(0);
while (!stk.empty()) {
int u = stk.top(); stk.pop();
max_dep = max(max_dep, dep[u]);
for (int nx : chld[u]) {
stk.push(nx);
dep[nx] = dep[u] + 1;
}
}
} ();
cerr << "blah\n\n";
for (int i = 0; i < n; i++) {
at_dep[dep[i]]++;
}
for (int i = 0; i < k; i++) process_whole_tree(i);
printf("%d %d\n", ans_tot, ans_moves);
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... |