arimattiの競プロ精進日記

競プロ問題の解法を記録していきます

ABC191 E - Come Back Quickly を解いてみた

問題リンク

atcoder.jp

問題概要

 N 個の町と町をつなぐ  M 本の道がある。道は一方通行であり、各道を通るのにかかるコストが与えられる(重み付き有向グラフ)。
与えられる道は出発する町=到着する町(自己ループ)もあり得るし、同じ町をつなぐ道が複数ある(多重辺)場合もある。
ある町から出発し、1 本以上の道を通りその町に戻ってくる際の最小コストを各町について求めよ。

制約

1 \leq N \leq 2000
1 \leq M \leq 2000

ポイント

  • 最短経路問題なのでダイクストラをしたい。
  • ただし、『1 本以上の道を通りその町に戻ってくる』には工夫が必要。
  • これを実現するために、町  i を出発した後、町  j を経由し、町  i に戻って来ることを考える。
  • 各町を始点としたダイクストラを行い、町  i → j と 町  j → i の時間を求める。
  •  i → j の最短時間を  dist [ i ][ j ] とすると、 i \neq j のすべての  j の中での  dist [ i ][ j ] + dist [ j ][ i ] の最小値を求める。自己ループがない場合はこれが答えである。
  • 自己ループがある場合を忘れてはいけない。上記求めた値と自己ループのコストとの min をとる。

感想

出発地点に戻ってくるパターンのダイクストラですね。これは覚えておきたい。各町の最小コストを求める & ある地点の経由と考えるとワーシャルフロイドが使えそう!と思いましたが  O(N ^ 3) で当然 TLE ですね。

コード

#include <bits/stdc++.h>
#include <atcoder/all>
#define rep(i,n) for (int i = 0; i < (n); i++)
using namespace std;
using namespace atcoder;
using ll = long long;
using P = pair<int, int>;
template<class T> inline bool chmin(T& a, T b) { if (a > b) {a = b; return true;} else return false;}
struct Edge {int to; ll w; Edge(int to, ll w) : to(to), w(w) {}; };
using Graph = vector<vector<Edge>>;
const ll INFL = 1LL<<60;

int main() {
  int n, m;
  cin >> n >> m;
  Graph G(n);
  vector<int> a(m), b(m);
  vector<ll> w(m);
  rep(i, m) {
    cin >> a[i] >> b[i] >> w[i];
    a[i]--; b[i]--;
    G[a[i]].push_back(Edge(b[i], w[i]));
  }

  vector<vector<ll>> dist(n, vector<ll>(n, INFL));
  // 出発地点を i としてダイクストラ法
  rep(i, n) {
    int s = i;
    dist[s][s] = 0;
    priority_queue< pair<ll, int>, vector<pair<ll, int>>, greater<pair<ll, int>> > que;
    que.push(make_pair(dist[s][s], s));
    while (!que.empty()) {
      int v = que.top().second;
      ll d = que.top().first;
      que.pop();
      if (d > dist[s][v]) continue;
      for (auto e : G[v]) {
        if (chmin(dist[s][e.to], dist[s][v]+e.w)) {
          que.push(make_pair(dist[s][e.to], e.to));
        }
      }
    }
    que = decltype(que)();
  }
  
  vector<ll> ans(n, INFL);
  rep(i, n) {
    rep(j, n) {
      if (i == j) continue;
      chmin(ans[i], dist[i][j]+dist[j][i]);
    }
  }
  // 自己ループがある場合は更新可能か調べる
  rep(i, m) {
    if (a[i] == b[i]) {
      chmin(ans[a[i]], w[i]);
    }
  }
  rep(i, n) {
    if (ans[i] == INFL) cout << -1 << endl;
    else cout << ans[i] << endl;    
  }
  return 0;
}