AtCoder Regular Contest 035 C - アットコーダー王国の交通事情

問題

 N個の都市がM本の道路で結ばれている。道路には長さがあり、各都市i,j間の最短距離の総和をSとする。初期状態からK本の新しい道路を建設していくとき、各道路が作られた後のSを求めよ。

解法

 新たな道路が都市X,Yを長さZで結ぶとき、各都市i,jについて最短距離を更新し得るのはi \to X\to Y \to jあるいはi \to Y \to X \to jのどちらか。つまり一般の都市aからbへのコストを保持しておけばこれはO(1)で求まる。i,jの選び方がそれぞれN通りあり、新たな道路はK本なので全体の計算量はO(KN^2)。初期状態において各都市間の最短距離を求めることについてもワーシャルフロイド法を用いればO(N^3)で計算できるので間に合う。

反省

 23分00秒(0WA)でAC。問題を理解できればとくに難しいところもなく。和を差分計算する必要があるかなと一瞬思ったが、計算量的には変わらないので面倒くさくて毎回O(N^2)かけて計算するままに。400^3が間に合うものなのかちょっと不安だったが495msで通った。つまり10^9でも簡単な処理ならギリギリ通ってしまう? このあたりの間に合うかどうかの感覚はよくわからないまま。

コード

#include"bits/stdc++.h"
using namespace std;
using ll = int64_t;

int main() {
    ll N, M;
    cin >> N >> M;
    
    struct Edge {
        ll to, cost;
    };
    vector<vector<Edge>> edges(N);
    for (ll i = 0; i < M; i++) {
        ll A, B, C;
        cin >> A >> B >> C;
        A--; B--;
        edges[A].push_back({ B, C });
        edges[B].push_back({ A, C });
    }

    ll K;
    cin >> K;
    vector<ll> X(K), Y(K), Z(K);
    for (ll i = 0; i < K; i++) {
        cin >> X[i] >> Y[i] >> Z[i];
        X[i]--; Y[i]--;
    }

    vector<vector<ll>> cost(N, vector<ll>(N, INT_MAX));
    for (ll i = 0; i < N; i++) {
        for (auto e : edges[i]) {
            cost[i][e.to] = min(cost[i][e.to], e.cost);
        }
        cost[i][i] = 0;
    }
    for (ll k = 0; k < N; k++) {
        for (ll i = 0; i < N; i++) {
            for (ll j = 0; j < N; j++) {
                cost[i][j] = min(cost[i][j], cost[i][k] + cost[k][j]);
            }
        }
    }

    auto getSum = [&]() {
        ll sum = 0;
        for (ll i = 0; i < N; i++) {
            for (ll j = i + 1; j < N; j++) {
                sum += cost[i][j];
            }
        }
        return sum;
    };

    for (ll i = 0; i < K; i++) {
        for (ll j = 0; j < N; j++) {
            for (ll k = 0; k < N; k++) {
                ll cost1 = cost[j][X[i]] + cost[Y[i]][k] + Z[i];
                ll cost2 = cost[j][Y[i]] + cost[X[i]][k] + Z[i];
                cost[j][k] = min({ cost[j][k], cost1, cost2 });
                cost[k][j] = min({ cost[k][j], cost1, cost2 });
            }
        }
        cost[X[i]][Y[i]] = min(cost[X[i]][Y[i]], Z[i]);
        cost[Y[i]][X[i]] = min(cost[Y[i]][X[i]], Z[i]);

        cout << getSum() << endl;
    }
}