Submission #1269340

#	Time	Username	Problem	Language	Result	Execution time	Memory
1269340		Davdav1232	Railway (BOI17_railway)	C++20	52 / 100	135 ms	32704 KiB

railway

#include <bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef pair<int,int> vii;

vector<pair<map<int,int>*, int>> things;
vector<vector<vii>> G;
vi par_num;
vi ans;
vvi colors;
vi S;
int k;

void dfs(int node, int p) {
    int d = 0;
    int m = -1;
    int j = -1; // initialize heavy-child index

    // find heavy child
    for (int i = 0; i < (int)G[node].size(); i++) {
        if (G[node][i].first == p) {
            par_num[node] = G[node][i].second;
            continue;
        }
        d++;
        dfs(G[node][i].first, node);

        // defensive assert: child map must exist
        assert(things[G[node][i].first].first != nullptr);

        if ((int)things[G[node][i].first].first->size() > m) {
            m = (int)things[G[node][i].first].first->size();
            j = i;
        }
    }

    if (d == 0) {
        // leaf: allocate new map
        things[node].first = new map<int,int>();
        for (int c : colors[node]) {
            (*things[node].first)[c]++;
        }
        if ((int)things[node].first->size() - things[node].second >= k && node != 0)
            ans.push_back(par_num[node]);
        return;
    }

    // if no heavy child chosen (defensive), pick first non-parent
    if (j == -1) {
        for (int i = 0; i < (int)G[node].size(); ++i) {
            if (G[node][i].first != p) { j = i; break; }
        }
    }

    // inherit heavy child's map
    things[node] = things[G[node][j].first];

    // add current node's colors
    for (int c : colors[node]) {
        auto &mp = *things[node].first;
        mp[c]++;
        if (mp[c] == S[c]) things[node].second++;
    }

    // merge other children
    for (int i = 0; i < (int)G[node].size(); ++i) {
        if (i == j || G[node][i].first == p) continue;
        int child = G[node][i].first;
        things[node].second += things[child].second;

        for (auto &kv : *things[child].first) {
            int color = kv.first;
            int count = kv.second;
            auto &mp = *things[node].first;
            int before = mp[color];
            mp[color] += count;
            if (before < S[color] && mp[color] >= S[color]) things[node].second++;
        }
    }

    if (node != 0 && (int)things[node].first->size() - things[node].second >= k)
        ans.push_back(par_num[node]);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m;
    cin >> n >> m >> k;

    colors.assign(n, {});
    G.assign(n, {});
    par_num.assign(n, 0);
    things.assign(n, {nullptr, 0});
    S.assign(m, 0);

    for (int i = 0; i < n-1; i++) {
        int a,b;
        cin >> a >> b;
        --a; --b;
        G[a].push_back({b, i+1});
        G[b].push_back({a, i+1});
    }

    for (int i = 0; i < m; i++) {
        int s; cin >> s;
        S[i] = s;
        for (int j = 0; j < s; j++) {
            int a; cin >> a; --a;
            colors[a].push_back(i);
        }
    }

    dfs(0, -1);

    sort(ans.begin(), ans.end());
    cout << ans.size() << "\n";
    for (int x : ans) cout << x << " ";
    cout << "\n";

    // cleanup: delete unique maps
    unordered_set<map<int,int>*> uniq;
    for (auto &p : things) if (p.first) uniq.insert(p.first);
    for (auto mp : uniq) delete mp;

    return 0;
}

• Subtask 1: 0 / 8
• Subtask 2: 0 / 15
• Subtask 3: 7 / 7
• Subtask 4: 29 / 29
• Subtask 5: 16 / 16
• Subtask 6: 0 / 25