제출 #1206626

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

Main

#include <bits/stdc++.h>
using namespace std;

constexpr int INF32 = 1e9;
struct Tree {
    int n;                     // number of nodes
    vector<vector<int>> g;     // children
    vector<int> lang;          // language of each node
    vector<int> tin, tout, eul, depth, sz, heavy;
    int timer = 0, maxDepth = 0;

    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) {                       // sizes, depths and 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;
    }
};

/* ────────────────────────────────────────────────────────────────────────── */
/* helpers shared by the two phases                                          */
struct CounterByDepth {
    vector<vector<int>> byDepth;   // tin-indices per depth (monotonically increasing)

    explicit CounterByDepth(int maxD) : byDepth(maxD + 1) {}

    void add(int d, int tin) { byDepth[d].push_back(tin); }

    // number of nodes at depth d inside Euler interval [l,r]
    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));
    }
};
/* ────────────────────────────────────────────────────────────────────────── */
/* global answer                                                             */
struct Best { long long team = -1; long long sw = INF32; };

inline void upd(Best &best, long long team, long long sw) {
    if (team > best.team || (team == best.team && sw < best.sw))
        best.team = team, best.sw = sw;
}
/* ────────────────────────────────────────────────────────────────────────── */
/* phase 1 : “rare” colours                                                  */
void solveRareColour(int c,
                     const vector<int>& nodes,            // all nodes of colour c
                     const unordered_map<int,int>& cap,   // cap[d] = #nodes of c at depth d
                     const vector<int>& tin,
                     const vector<int>& tout,
                     const CounterByDepth& counter,
                     const Tree& T,
                     Best& best)
{
    vector<int> depthList; depthList.reserve(cap.size());
    for (auto &kv : cap) depthList.push_back(kv.first);

    for (int u : nodes) {
        long long val = 0;
        for (int d : depthList) {
            int inside = counter.count(d, tin[u], tout[u]);
            int limit  = cap.at(d);
            val += min(inside, limit);
        }
        // initial good employees of colour c in the subtree:
        long long good = 0;
        for (int w : nodes)
            if (tin[w] >= tin[u] && tin[w] <= tout[u]) ++good;

        upd(best, val, val - good);
    }
}

/* ────────────────────────────────────────────────────────────────────────── */
/* phase 2 : “heavy” colours – one DSU-on-tree run per colour                */
struct HeavySolver {
    const Tree& T;
    const vector<int>& cap;          // cap[d] for current colour
    Best& best;
    int curColour;
    vector<int> freq;                // current freq at depth d
    long long curSum = 0;            // Σ min(freq[d],cap[d])
    long long goodCnt = 0;           // employees of colour c in current subtree

    HeavySolver(const Tree& T_, const vector<int>& cap_, int colour, Best& b)
        : T(T_), cap(cap_), best(b), curColour(colour),
          freq(T_.maxDepth + 1, 0) {}

    void add(int node) {
        int d = T.depth[node];
        int old = freq[d]++;
        curSum += min(old+1, cap[d]) - min(old, cap[d]);
        if (T.lang[node] == curColour) ++goodCnt;
    }
    void remove(int node) {
        int d = T.depth[node];
        int old = freq[d]--;
        curSum += min(old-1, cap[d]) - min(old, cap[d]);
        if (T.lang[node] == curColour) --goodCnt;
    }

    void addSubtree(int l, int r) {          // inclusive Euler interval
        for (int i = l; i <= r; ++i) add(T.eul[i]);
    }
    void remSubtree(int l, int r) {
        for (int i = l; i <= r; ++i) remove(T.eul[i]);
    }

    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]);

        add(u);

        if (T.lang[u] == curColour)
            upd(best, curSum, curSum - goodCnt);

        if (!keep) remSubtree(T.tin[u], T.tout[u]);
    }
};

void solveHeavyColour(int colour,
                      const vector<int>& nodes,
                      const CounterByDepth& counter,
                      Tree& T,
                      Best& best)
{
    vector<int> cap(T.maxDepth + 1, 0);
    for (int v : nodes) ++cap[T.depth[v]];

    HeavySolver S(T, cap, colour, best);
    S.dfs(0, false);
}
/* ────────────────────────────────────────────────────────────────────────── */
/* main driver                                                               */
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, b; v < N; ++v) {          // bosses
        cin >> b;
        T.addEdge(b, v);
    }

    T.dfs0(0);                                // preprocessing

    /* depth lists for O(log N) counting */
    CounterByDepth counter(T.maxDepth);
    for (int v = 0; v < N; ++v)
        counter.add(T.depth[v], T.tin[v]);

    /* bucket nodes by colour */
    vector<vector<int>> byColour(K);
    for (int v = 0; v < N; ++v) byColour[T.lang[v]].push_back(v);

    const int B = (int)std::sqrt(N) + 1;       // threshold between rare / heavy
    Best best;                                 // final answer

    /* phase 1 – rare colours */
    for (int c = 0; c < K; ++c) {
        int cnt = (int)byColour[c].size();
        if (cnt == 0 || cnt >= B) continue;

        unordered_map<int,int> cap;
        for (int v : byColour[c]) ++cap[T.depth[v]];

        solveRareColour(c, byColour[c], cap,
                        T.tin, T.tout, counter, T, best);
    }

    /* phase 2 – heavy colours */
    for (int c = 0; c < K; ++c) {
        int cnt = (int)byColour[c].size();
        if (cnt < B) continue;                 // handled already

        solveHeavyColour(c, byColour[c], counter, T, best);
    }

    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