IOI '11 P1 - Tropical Garden (Standard I/O)

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>

using namespace std;

const int INF = 1e9 + 10;

int main() {
    int n, m, p;
    cin >> n >> m >> p;
    vector<vector<pair<int, int>>> g(n);
    for (int i = 0; i < m; ++i) {
        int u, v; cin >> u >> v;
        g[u].emplace_back(i, v);
        g[v].emplace_back(i, u);
    }

    // split each node into two nodes:
    //  - can exit through smallest edge
    //  - cannot exit through smallest edge
    vector<int> nxt(2*n);
    for (int u = 0; u < n; ++u) {
        sort(g[u].begin(), g[u].end());
        nxt[2*u] = g[u][0].second; // will exit through smallest edge
        nxt[2*u+1] = g[u][g[u].size() == 1 ? 0 : 1].second; // will exit through second smallest edge (if exists)
    }
    for (int u = 0; u < 2*n; ++u) {
        nxt[u] <<= 1;
        // entered nxt[u] through its smallest edge?
        if (g[nxt[u]>>1][0].second == (u>>1)) ++nxt[u];
    }

    // there are now two target nodes: 2*p and 2*p+1
    int targetCycle[2]{INF, INF};
    for (int target = 2*p; target <= 2*p+1; ++target) {
        vector<bool> seen(2*n);
        int cur = target, cycleSize = 0;
        while (!seen[cur]) {
            seen[cur] = true;
            ++cycleSize;
            cur = nxt[cur];
        }
        if (cur == target) targetCycle[target-2*p] = cycleSize;
    }

    vector<vector<int>> prv(2*n);
    for (int u = 0; u < 2*n; ++u) prv[nxt[u]].push_back(u);

    // reverse bfs from each target to find distances
    vector<int> targetDistance[2];
    for (int target = 2*p; target <= 2*p+1; ++target) {
        int id = target-2*p;
        targetDistance[id] = vector<int>(2*n, INF);
        queue<int> q;
        q.push(target);
        targetDistance[id][target] = 0;
        while (!q.empty()) {
            int v = q.front();
            q.pop();
            for (int u : prv[v]) {
                if (targetDistance[id][u] == INF) {
                    targetDistance[id][u] = targetDistance[id][v] + 1;
                    q.push(u);
                }
            }
        }
    }

    int q; cin >> q;
    while (q--) {
        int k, ans = 0; cin >> k;
        for (int i = 0; i < n; ++i) {
            int u = 2*i; // must start in node that can exit through smallest
            bool can = false;
            // check to see if it can reach either target node in k moves
            for (int target = 2*p; target <= 2*p+1 && !can; ++target) {
                int id = target-2*p;
                if (targetDistance[id][u] > k) continue;
                can = (k - targetDistance[id][u]) % targetCycle[id] == 0;
            }
            ans += can;
        }
        cout << ans << ' ';
    }
    cout << endl;
    return 0;
}