AOJ 0627 Train Fare (鉄道運賃)

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0627

解法

まず首都からの最短経路を求めるほかどうしようもないので，BFSで最初に求めておく．
このとき，任意の頂点について，すべての最短経路でのその頂点の親と子を求めておく．
経路復元するときと同じ感じでやればいい．

クエリが投げられたら，その辺が最短経路で使われているか確認する．
これは辺 (u, v) に対して u と v が親子関係か調べれば良い．
使われていれば，仮に親を u，子を v とすると，u の子どもの集合から v を取り除き，v の親の集合から u を取り除く．
親の集合のサイズが 0 になったときが，最短経路が崩れてしまう条件なので，この時 v の親の集合が 0 になったら，v を親として，さっきと同じことを伝播させてやっていく．
取り除いた頂点は別に管理しておけば良く，クエリの最後にはこの取り除いた頂点のサイズを答える．

ソースコード

#include <bits/stdc++.h>
using namespace std;
 
using pii = pair<int, int>;

constexpr int INF = 1e9;

void erase(int p, int ch, vector<set<int>>& par, vector<set<int>>& child, set<int>& upset) {
    queue<pii> que;
    que.push(make_pair(p, ch));
    while(!que.empty()) {
        int pp = que.front().first;
        int c = que.front().second;
        que.pop();
        par[c].erase(pp);
        child[pp].erase(c);
        if(par[c].empty() && upset.count(c) == 0) {
            upset.insert(c);
            for(auto& cc : child[c]) {
                que.push(make_pair(c, cc));
            }
        }
    }
}

int main() {
    int N, M, Q;
    cin >> N >> M >> Q;
    vector<vector<int>> g(N);
    vector<pii> es(M);
    for(int i = 0; i < M; ++i) {
        int u, v;
        cin >> u >> v;
        u--; v--;
        es[i].first = u;
        es[i].second = v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    vector<int> d(N, INF);
    vector<set<int>> par(N);
    vector<set<int>> child(N);
    queue<pii> que;
    que.push(make_pair(0, 0));
    d[0] = 0;
    while(!que.empty()) {
        int v = que.front().first;
        int cost = que.front().second;
        que.pop();
        for(auto to : g[v]) {
            if(d[to] >= cost + 1) {
                child[v].insert(to);
                par[to].insert(v);
            }
            if(d[to] > cost + 1) {
                d[to] = cost + 1;
                que.push(make_pair(to, d[to]));
            }
        }
    }

    set<int> upset;
    while(Q--) {
        int r;
        cin >> r;
        r--;
        int p = es[r].first, c = es[r].second;
        if(d[p] + 1 != d[c]) {
            swap(p, c);
        }
        if(d[p] + 1 == d[c]) {
            erase(p, c, par, child, upset);
        }
        cout << upset.size() << endl;
    }
}

感想

もっと簡単に書ける気がするんだけど，実装方針あんまり考えずに実装してしまった．