Submission #536530

#	Time	Username	Problem	Language	Result	Execution time	Memory
536530		Stickfish	Railway (BOI17_railway)	C++17	100 / 100	290 ms	41900 KiB

railway

#include <iostream>
#include <algorithm>
#include <vector>
#include <map>
using namespace std;

const int MAXN = 1e5 + 123;
vector<int> edg[MAXN];
vector<int> redg[MAXN];
int rtnum[MAXN];
int depth[MAXN];
int timer = 0;
int tin[MAXN];
int tout[MAXN];
int to_root_add[MAXN];

pair<int, int> ghash(int v, int u) {
    if (v > u)
        return {u, v};
    return {v, u};
}

void dfsinit(int v, map<pair<int, int>, int>& mp) {
    tin[v] = timer++;
    if (redg[v].size()) {
        for (int j = 0;; ++j) { 
            int u = redg[v][j];
            if (j >= redg[u].size())
                break;
            redg[v].push_back(redg[u][j]);
        }
    }
    for (auto u : edg[v]) {
        if (redg[v].size() && u == redg[v][0])
            continue;
        rtnum[u] = mp[ghash(v, u)];
        redg[u].push_back(v);
        depth[u] = depth[v] + 1;
        dfsinit(u, mp);
    }
    tout[v] = timer;
}

int lca(int u, int v) {
    if (depth[u] > depth[v]) {
        int t = depth[u] - depth[v];
        for (int j = redg[u].size(); j >= 0; --j) {
            if (t & (1 << j)) {
                u = redg[u][j];
                t -= (1 << j);
            }
        }
    }
    if (depth[v] > depth[u])
        return lca(v, u);
    if (u == v)
        return v;
    //cout << "MYUKU" << ' ' << u << ' ' << v << ": " << depth[u] << ' ' << depth[v] << endl;
    for (int j = redg[u].size() - 1; j >= 0; --j) {
        if (j < redg[u].size() && redg[u][j] != redg[v][j]) {
            u = redg[u][j];
            v = redg[v][j];
        }
    }
    //cout << "UNYU~" << endl;
    return redg[u][0];
}

bool cmptin(int i, int j) {
    return tin[i] < tin[j];
}

int get_ans_dfs(int v, int k, vector<int>& ans) {
    int cnt = 0;
    for (auto u : edg[v]) {
        if (redg[v].size() && u == redg[v][0])
            continue;
        cnt += get_ans_dfs(u, k, ans);
    }
    cnt += to_root_add[v];
    //cout << v << ": " << cnt << endl;
    //if (cnt > 1)
        //cout << "ERROR" << endl;
    if (cnt >= k)
        ans.push_back(rtnum[v]);
    return cnt;
}

signed main() {
    int n, m, k;
    cin >> n >> m >> k;
    map<pair<int, int>, int> mp;
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        --u, --v;
        edg[u].push_back(v);
        edg[v].push_back(u);
        mp[ghash(u, v)] = i;
    }
    dfsinit(0, mp);
    while (m--) {
        int s;
        cin >> s;
        vector<int> vc(s);
        for (int i = 0; i < s; ++i) {
            cin >> vc[i];
            --vc[i];
        }
        sort(vc.begin(), vc.end(), cmptin);
        //for (auto x : vc)
            //cout << x << ' ';
        //cout << endl;
        int a0 = lca(vc[0], vc[1]);
        ++to_root_add[vc[0]];
        ++to_root_add[vc[1]];
        to_root_add[a0] -= 2;
        for (int i = 2; i < s; ++i) {
            int ai = lca(vc[i - 1], vc[i]);
            if (depth[ai] >= depth[a0]) {
                ++to_root_add[vc[i]];
                --to_root_add[ai];
            } else {
                ++to_root_add[a0];
                ++to_root_add[vc[i]];
                to_root_add[ai] -= 2;
                a0 = ai;
            }
        }
    }
    vector<int> ans;
    get_ans_dfs(0, k, ans);
    sort(ans.begin(), ans.end());
    cout << ans.size() << '\n';
    for (auto x : ans)
        cout << x << ' ';
    cout << '\n';
}

Compilation message (stderr)

railway.cpp: In function 'void dfsinit(int, std::map<std::pair<int, int>, int>&)':
railway.cpp:28:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   28 |             if (j >= redg[u].size())
      |                 ~~^~~~~~~~~~~~~~~~~
railway.cpp: In function 'int lca(int, int)':
railway.cpp:60:15: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   60 |         if (j < redg[u].size() && redg[u][j] != redg[v][j]) {
      |             ~~^~~~~~~~~~~~~~~~

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