제출 #1206640

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1206640		ocean	Team Coding (EGOI24_teamcoding)	C++17	31 / 100	47 ms	29508 KiB

Main

/**************************************************************
 *  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;
}

• Subtask 1: 12 / 12
• Subtask 2: 19 / 19
• Subtask 3: 0 / 27
• Subtask 4: 0 / 23
• Subtask 5: 0 / 19