Codeforces Round #Pi (Div. 2) E. President and Roads

最初はすぎむさんの解法だけをお借りするつもりだったのですが,よくわからないTLEに悩まされたので結局ソースコードまでお借りしました。

問題

Problem - E - Codeforces

解法

s から t までの最短経路は普通は複数通りあります。で,問題の解答としては,「最短経路を通る際この辺は絶対使う」という場合は "YES" を,「最短経路を通る際この辺を使う可能性がある,かつ辺を通るのにかかる時間が 1 以上」だったら "CAN 1" を,また,その他の場合は,その辺が u -> v という辺であったとして, (s から u までの距離) + (u から v までの距離) + (v から t までの距離) が s から t までの距離より小さくなりうるなら "CAN ほげほげ" となり,そうでないなら "NO" を返せば OK ,みたいな感じです。

このなかで一番処理が難しいのがある辺について,「この辺は絶対使う」となるかどうかの判定です。残りはこれの説明をします。

まずは最短経路として使う可能性がある辺を列挙しています。ソースコード中では dijkstra のあとで main 関数の中でやってることで, (s から u までの距離) + (u から v までの距離) + (v から t までの距離) が s から t までの距離と等しくなるなら使う可能性があります。

このようにして集めた辺について,その辺を使う「確率」を計算します。(_dijkstra でやっていること。普通にdoubleで計算しても多分大丈夫だが下のコードではとするやつでやっている)この確率が1ならその辺は絶対に使うということなので "YES" を返せばよく, そうでないなら別の処理をする,という感じです。

ll MOD = 1e9+7;
struct edge {
    int v;
    ll w;
    edge() {}
    edge(int v, ll w) : v(v), w(w) {};
};

// extgcd
ll extgcd(ll a, ll b, ll& x, ll& y) {
    ll d = a;
    if (b != 0) {
        d = extgcd(b, a % b, y, x);
        y -= (a / b) * x;
    } else {
        x = 1; y = 0;
    }
    return d;
}

// mod_inverse
ll mod_inverse(ll a, ll m) {
    ll x, y;
    extgcd(a, m, x, y);
    return (m+x%m) % m;
}

vector<ll> dijkstra(int n, vector<vector<edge> >& G, int s) {
    vector<ll> d(n, LLONG_MAX/10); d[s] = 0;
    priority_queue<pair<ll, int> > que;
    que.push(make_pair(0ll, s));
    while (!que.empty()) {
        auto p = que.top(); que.pop();
        int u = p.second;
        ll dist = -p.first;
        if (dist > d[u]) continue;
        for (edge e : G[u]) {
            if (d[e.v] > d[u]+e.w) {
                d[e.v] = d[u] + e.w;
                que.push(make_pair(-d[e.v], e.v));
            }
        }
    }
    return d;
}

vector<ll> _dijkstra(int n, vector<vector<edge> >& G, int s) {
    vector<ll> d(n, LLONG_MAX/10); d[s] = 0;
    vector<ll> x(n); x[s] = 1;
    priority_queue<pair<ll, int> > que;
    que.push(make_pair(0ll, s));
    while (!que.empty()) {
        auto p = que.top(); que.pop();
        int u = p.second;
        ll dist = -p.first;
        if (dist > d[u]) continue;
        for (edge e : G[u]) {
            if (d[e.v] > d[u]+e.w) {
                d[e.v] = d[u] + e.w;
                que.push(make_pair(-d[e.v], e.v));
            }
            x[e.v] = (x[e.v] + x[u] * mod_inverse(G[u].size(), MOD)) % MOD;
        }
    }
    return x;
}

int main() {
    int n, m, s, t;
    cin >> n >> m >> s >> t;
    s--; t--;
    vector<int> a(m), b(m), l(m);
    vector<vector<edge> > G(n), rG(n);
    for (int i = 0; i < m; i++) {
        scanf("%d %d %d", &a[i], &b[i], &l[i]);
        a[i]--; b[i]--;
        G[a[i]].emplace_back(b[i], l[i]);
        rG[b[i]].emplace_back(a[i], l[i]);
    }
    vector<ll> d = dijkstra(n, G, s), rd = dijkstra(n, rG, t);
    ll D = d[t];
    vector<vector<edge> > g(n);
    for (int i = 0; i < m; i++) {
        if (d[a[i]] + rd[b[i]] + l[i] == D) g[a[i]].emplace_back(b[i], l[i]);
    }
    vector<ll> x = _dijkstra(n, g, s);
    for (int i = 0; i < m; i++) {
        if (g[a[i]].size() == 1 && d[a[i]] + rd[b[i]] + l[i] == D && x[a[i]] == 1 && x[b[i]] == 1) printf("YES\n");
        else if (d[a[i]] + rd[b[i]] + 1 >= D) printf("NO\n");
        else printf("CAN %I64d\n", (d[a[i]] + rd[b[i]] + l[i]) - (D-1));
    }
}