Submission #672662

#	Time	Username	Problem	Language	Result	Execution time	Memory
672662		bogdanvladmihai	Railway (BOI17_railway)	C++17	100 / 100	330 ms	47304 KiB

railway

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

int N, M, K;
vector<vector<int>> tree, g;
map<pair<int, int>, int> id;

vector<int> pref;
vector<int> st, en, lvl;
vector<vector<int>> up;

set<int> answer;

void dfs(int u, int dad = 0) {
    static int time = 0;
    st[u] = ++time;

    up[0][u] = dad;
    for (int l = 1; (1 << l) <= N; l++) {
        up[l][u] = up[l - 1][up[l - 1][u]];
    }

    for (const int &v : tree[u]) {
        if (v != dad) {
            lvl[v] = lvl[u] + 1;
            dfs(v, u);
        }
    }

    en[u] = time;
}

bool upper(int u, int v) {
    return (bool)(st[u] <= st[v] && en[v] <= en[u]);
}

int lca(int u, int v) {
    if (upper(u, v)) {
        return u;
    }
    if (upper(v, u)) {
        return v;
    }

    for (int l = ceil(log2(N)); l >= 0; l--) {
        if (!upper(up[l][u], v)) {
            u = up[l][u];
        }
    }

    return up[0][u];
}

int getVirtualTree(vector<int> &ver) {
    sort(ver.begin(), ver.end(), [&](int i, int j) {
        return st[i] < st[j];
    });

    int size = (int)ver.size();
    for (int i = 0; i < size - 1; i++) {
        ver.push_back(lca(ver[i], ver[i + 1]));
    }
    sort(ver.begin(), ver.end(), [&](int i, int j) {
        return st[i] < st[j];
    });
    ver.erase(unique(ver.begin(), ver.end()), ver.end());

    for (const int &u : ver) {
        g[u].clear();
    }

    vector<int> s;
    for (const int &u : ver) {
        while ((int)s.size() > 1 && !upper(s.back(), u)) {
            g[s[(int)s.size() - 2]].push_back(s.back());
            s.pop_back();
        }

        s.push_back(u);
    }

    while ((int)s.size() > 1) {
        g[s[(int)s.size() - 2]].push_back(s.back());
        s.pop_back();
    }

    return s[0];
}

void dfsVirtual(int u) {
    for (const int &v : g[u]) {
        pref[u]--;
        pref[v]++;

        dfsVirtual(v);
    }
}

void dfsSums(int u, int dad = -1) {
    for (const int &v : tree[u]) {
        if (v != dad) {
            dfsSums(v, u);

            if (pref[v] >= K) {
                answer.insert(id[make_pair(u, v)]);
            }
            pref[u] += pref[v];
        }
    }
}

int main() {
    cin >> N >> M >> K;

    tree.resize(N);
    g.resize(N);
    for (int i = 0; i < N - 1; i++) {
        int a, b; cin >> a >> b;
        a--; b--;

        tree[a].push_back(b);
        tree[b].push_back(a);
        id[make_pair(a, b)] = id[make_pair(b, a)] = i + 1;
    }

    st.resize(N);
    en.resize(N);
    lvl.resize(N);
    up.resize(ceil(log2(N)) + 1, vector<int>(N));
    lvl[0] = 1;
    dfs(0);

    pref.resize(N);

    for (int i = 0; i < M; i++) {
        int V; cin >> V;
        vector<int> ver(V);
        for (int &u : ver) {
            cin >> u;
            u--;
        }

        int virtualRoot = getVirtualTree(ver);
        dfsVirtual(virtualRoot);
    }

    dfsSums(0);

    cout << (int)answer.size() << "\n";
    for (const int &edge : answer) {
        cout << edge << " ";
    }
    cout << "\n";

    return 0;
}

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