mayoko’s diary

プロコンとかいろいろ。

Codeforces Round #372 (Div. 1) B. Complete The Graph

いやなんか時間とってゆっくり競プロやりたいんですけど…

問題

codeforces.com

コスト付きの無向グラフが与えられる。いくつかの辺のコストは好きに選べるので, s から t への最短経路の大きさを L にしたい。このようなことが可能であれば, そのうちの 1 つを答えよ。

解法

まず自明に解がない場合を省きます。というのは, いじれるすべての辺のコストを INF にしても最短経路が L 未満の場合と, いじれるすべての辺のコストを 1 にしても最短経路が L より大きい場合です。

逆に, これ以外の場合は, 解が存在することが示せます。これは

  • いじれるすべての辺のコストを 1 にすると最短経路は L 未満, 全て INF にすると L より大きい
  • いじれる辺のいずれかのコストを 1 増やすと最短経路は 1 増えるかそのまま

から言えます。

で, どう具体的にアルゴリズムに落とすかというと,

  • いじれる辺のすべてのコストを 1 にセット
  • 以下を繰り返す
    • dijkstra して t までの距離を求める。L なら終了
    • そうでない場合, s から t までの最短経路にいじれる辺があるはず。その辺のコストを L-d[t]+1 にする。

というようにやれば, いじる辺は M 本しかないので O(M^2 log N) になります(これでギリギリ通った)。さらに, 最初にすべてのコストを 1 にしたときの最短経路上の辺のみを使うようにすると O(NMlog N) になって高速になります。

struct edge {
    int v;
    ll w;
    int id;
    edge() {}
    edge(int v, ll w, int id) : v(v), w(w), id(id) {};
};

vector<edge> G[1010];

vector<ll> dijkstra(int n, int s, int t) {
    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;
        if (u == t) return d;
        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;
}

void no() {
    cout << "NO" << endl;
    exit(0);
}

const int MAXM = 10101;
int U[MAXM], V[MAXM], W[MAXM];
bool use[MAXM];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N, M, L, s, t;
    cin >> N >> M >> L >> s >> t;
    const ll INF = 1e9;
    for (int i = 0; i < M; i++)
        cin >> U[i] >> V[i] >> W[i];
    for (int i = 0; i < M; i++) if (W[i]) {
        G[U[i]].emplace_back(V[i], W[i], i);
        G[V[i]].emplace_back(U[i], W[i], i);
    }
    auto d = dijkstra(N, s, t);
    if (d[t] < L) no();
    if (d[t] == L) {
        cout << "YES" << endl;
        for (int i = 0; i < M; i++) {
            ll w = (W[i]==0) ? INF : W[i];
            cout << U[i] << " " << V[i] << " " << w << endl;
        }
        return 0;
    }
    for (int i = 0; i < M; i++) if (!W[i]) {
        G[U[i]].emplace_back(V[i], 1, i);
        G[V[i]].emplace_back(U[i], 1, i);
    }
    d = dijkstra(N, s, t);
    if (d[t] > L) no();
    {
        int now = t;
        while (now != s) {
            for (edge e : G[now]) {
                if (d[e.v] + e.w == d[now]) {
                    use[e.id] = true;
                    now = e.v;
                    break;
                }
            }
        }
    }
    while (1) {
        for (int i = 0; i < N; i++)
            G[i].clear();
        for (int i = 0; i < M; i++) {
            ll w = INF;
            if (W[i] == 0 && use[i]) w = 1;
            if (W[i]) w = W[i];
            G[U[i]].emplace_back(V[i], w, i);
            G[V[i]].emplace_back(U[i], w, i);
        }
        auto d = dijkstra(N, s, t);
        if (d[t] == L) break;
        int now = t;
        while (now != s) {
            int id = -1;
            for (edge e : G[now]) {
                if (d[e.v]+e.w == d[now]) {
                    id = e.id;
                    now = e.v;
                    break;
                }
            }
            if (W[id] == 0) {
                W[id] = L-d[t]+1;
                break;
            }
        }
    }
    cout << "YES" << endl;
    for (int i = 0; i < M; i++) {
        ll w = INF;
        if (W[i] == 0 && use[i]) w = 1;
        if (W[i]) w = W[i];
        cout << U[i] << " " << V[i] << " " << w << endl;
    }
    return 0;
}