Submission #1146132

#TimeUsernameProblemLanguageResultExecution timeMemory
1146132monkey133Team Coding (EGOI24_teamcoding)C++20
100 / 100
241 ms31420 KiB
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
 
#define X first
#define Y second
#define pb push_back
#define ALL(v) v.begin(), v.end()
#define SZ(a) ((int)a.size())
#ifdef bbq
#include <experimental/iterator>
#define safe cerr<<__PRETTY_FUNCTION__<<" line "<<__LINE__<<" safe\n"
#define debug(a...) debug_(#a, a)
#define orange(a...) orange_(#a, a)
void debug_(auto s, auto ...a) {
    cerr << "\e[1;32m(" << s << ") = (";
    int f = 0;
    (..., (cerr << (f++ ? ", " : "") << a));
    cerr << ")\e[0m\n";
}
void orange_(auto s, auto L, auto R) {
    cerr << "\e[1;33m[" << s << "] = [";
    using namespace experimental;
    copy(L, R, make_ostream_joiner(cerr, ", "));
    cerr << "]\e[0m\n";
}
#else
#define safe ((void)0)
#define debug(...) safe
#define orange(...) safe
#endif
 
const int K = 320;
const int MAXN = 100005;
 
// Global arrays and vectors
vector<int> G[MAXN], stk, dep, in, out;
int dft;
 
// Standard DFS for Euler tour; pushes positive numbers for entry and negative for exit.
void dfs(int u, int d) {
    dep[u] = d, in[u] = ++dft;
    stk.pb(u + 1); // push entry marker (1-indexed)
    for (int v : G[u])
        dfs(v, d + 1);
    stk.pb(-u - 1); // push exit marker
    out[u] = dft;
}
 
// Returns true if u is an ancestor of v.
bool ancestor(int u, int v) {
    return in[u] <= in[v] && out[u] >= out[v];
}
 
// Overloaded operator+ for pii, in case you need it.
pii operator+(const pii &a, const pii &b) {
    return pii(a.X + b.X, a.Y + b.Y);
}
 
// Main
int main(){
    ios::sync_with_stdio(0), cin.tie(0);
    
    int n, k;
    cin >> n >> k;
    vector<int> arr(n), pa(n, -1);
    vector<vector<int>> buk(k), layer(k);
    dep.resize(n), in.resize(n), out.resize(n);
    
    // Read colors and bucket them by color.
    for (int i = 0; i < n; ++i) {
        cin >> arr[i];
        buk[arr[i]].pb(i);
    }
    // Read parent array and build the tree.
    for (int i = 1; i < n; ++i) {
        cin >> pa[i];
        G[pa[i]].pb(i);
    }
    dfs(0, 0);
    
    // Our answer is stored as (res, -swaps) so that lexicographical max works.
    pii ans = pii(0, 0);
 
    // 'proc' updates the candidate answer using given parameters.
    auto proc = [&](pii &cur, int f, int l, int a) {
        debug(f, l, a);
        assert(f >= a && l >= a);
        cur.X += min(f, l);
        cur.Y -= min(f, l) - a;
    };
 
    // 'up' merges the vector of node i directly into its parent's vector.
    auto up = [&](auto &len, int i) {
        if (pa[i] != -1) {
            // Always merge the smaller vector into the larger one.
            if (SZ(len[pa[i]]) < SZ(len[i]))
                len[pa[i]].swap(len[i]);
            for (int j = 0; j < SZ(len[i]); ++j)
                len[pa[i]][SZ(len[pa[i]]) - j - 1] = 
                    len[i][SZ(len[i]) - j - 1] + 
                    len[pa[i]][SZ(len[pa[i]]) - j - 1];
        }
    };
 
    // 'solve_big' processes colors with many occurrences (big colors).
    auto solve_big = [&](int c) {
        vector<vector<pii>> len(n); // len[i] is a vector for node i
        vector<int> f(n);
        for (int i : buk[c])
            ++f[dep[i]];
        int vis = 0;
        for (int x : stk) {
            int flag = (arr[abs(x) - 1] == c);
            if (x > 0) vis += flag;
            else {
                int i = -x - 1; 
                vis -= flag;
                len[i].pb(pii(1, flag));
                if (flag) {
                    if (vis == 0) {
                        pii cur = pii(0, 0);
                        for (int j = 0; j < SZ(len[i]); ++j) {
                            int realf = f[dep[i] + j];
                            auto [reall, reala] = len[i][SZ(len[i]) - j - 1];
                            proc(cur, realf, reall, reala);
                        }
                        ans = max(ans, cur);
                    }
                }
                up(len, i);
            }
        }
    };
 
    // 'solve_small' processes colors with few occurrences.
    vector<int> small, visArr(n, -1), cnt(n), ina(n);
    auto solve_small = [&]() {
        vector<vector<int>> len(n);
        for (int c : small) {
            for (int i : buk[c]) {
                if (visArr[dep[i]] != c) {
                    visArr[dep[i]] = c;
                    layer[c].pb(dep[i]);
                }
            }
            sort(ALL(layer[c]));
        }
        for (int x : stk) {
            if (x > 0) continue;
            int i = -x - 1;
            len[i].pb(1);
            debug("exit", i);
            orange(ALL(len[i]));
            if (SZ(buk[arr[i]]) < K) {
                pii cur = pii(0, 0);
                for (int j : layer[arr[i]])
                    cnt[j] = ina[j] = 0;
                for (int j : buk[arr[i]]) {
                    ++cnt[dep[j]];
                    if (ancestor(i, j))
                        ++ina[dep[j]];
                }
                for (int j : layer[arr[i]])
                    if (j >= dep[i]) {
                        if (SZ(len[i]) - 1 >= j - dep[i])
                            proc(cur, cnt[j], len[i][SZ(len[i]) - (j - dep[i]) - 1], ina[j]);
                    }
                ans = max(ans, cur);
            }
            up(len, i);
            if (i != 0) {
                debug("done up", i);
                orange(ALL(len[pa[i]]));
            }
        }
    };
 
    // Process each color: if bucket size >= K then it's a big color; otherwise, mark it as small.
    for (int c = 0; c < k; ++c) {
        if (SZ(buk[c]) >= K)
            solve_big(c);
        else
            small.pb(c);
    }
    solve_small();
    cout << ans.X << " " << -ans.Y << "\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...