Submission #547328

#	Time	Username	Problem	Language	Result	Execution time	Memory
547328		Jomnoi	Railway (BOI17_railway)	C++17	100 / 100	170 ms	28040 KiB

railway

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

const int MAX_N = 1e5 + 10;

int timer, k;
vector <pair <int, int>> graph[MAX_N];
int depth[MAX_N], parent[MAX_N][20];
int st[MAX_N], pos[MAX_N], edge[MAX_N];
int qs[MAX_N];
vector <int> ans;

void dfs(const int &u, const int &p) {
    timer++;
    st[u] = timer;
    pos[timer] = u;

    for(int i = 1; i < 20; i++) {
        parent[u][i] = parent[parent[u][i - 1]][i - 1];
    }
    for(auto [v, i] : graph[u]) {
        if(v == p) {
            continue;
        }

        edge[v] = i;
        depth[v] = depth[u] + 1;
        parent[v][0] = u;
        dfs(v, u);
    }
}

int find_lca(const int &a, const int &b) {
    int u = a, v = b;
    if(depth[u] < depth[v]) {
        swap(u, v);
    }
    for(int i = 19; i >= 0; i--) {
        if(depth[parent[u][i]] >= depth[v]) {
            u = parent[u][i];
        }
    }
    for(int i = 19; i >= 0; i--) {
        if(parent[u][i] != parent[v][i]) {
            u = parent[u][i];
            v = parent[v][i];
        }
    }
    return (u == v ? u : parent[u][0]);
}

int dfs2(const int &u, const int &p) {
    int sum = qs[u];
    for(auto [v, i] : graph[u]) {
        if(v != p) {
            sum += dfs2(v, u);
        }
    }

    if(sum >= k) {
        ans.push_back(edge[u]);
    }
    return sum;
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    int n, m;
    cin >> n >> m >> k;
    for(int i = 1; i <= n - 1; i++) {
        int a, b;
        cin >> a >> b;
        graph[a].emplace_back(b, i);
        graph[b].emplace_back(a, i);
    }

    depth[0] = -1;
    dfs(1, -1);

    while(m--) {
        int s;
        cin >> s;
        vector <int> vec;
        while(s--) {
            int x;
            cin >> x;
            vec.push_back(x);
        }

        sort(vec.begin(), vec.end(), [&](const int &a, const int &b) {
            return st[a] < st[b];
        });

        stack <int> latest_lca;
        stack <int> node;
        node.push(vec[0]);
        for(int i = 1; i < vec.size(); i++) {
            while(!latest_lca.empty() and st[latest_lca.top()] >= st[find_lca(vec[i], node.top())]) {
                qs[node.top()]++;
                node.pop();
                qs[node.top()]++;
                node.pop();
                qs[latest_lca.top()] -= 2;
                node.push(latest_lca.top());
                latest_lca.pop();
            }
            latest_lca.push(find_lca(vec[i], node.top()));
            node.push(vec[i]);
        }
        while(!latest_lca.empty()) {
            qs[node.top()]++;
            node.pop();
            qs[node.top()]++;
            node.pop();
            qs[latest_lca.top()] -= 2;
            node.push(latest_lca.top());
            latest_lca.pop();
        }
    }

    dfs2(1, -1);

    cout << ans.size() << '\n';
    sort(ans.begin(), ans.end());
    for(auto v : ans) {
        cout << v << ' ';
    }
    return 0;
}

Compilation message (stderr)

railway.cpp: In function 'int main()':
railway.cpp:98:26: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   98 |         for(int i = 1; i < vec.size(); i++) {
      |                        ~~^~~~~~~~~~~~

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