最短経路問題をたとえTLE（時間オーバー）しようが色々な解法で解く

概要：想定解で単純に解ける問題を、敢えてワーシャルフロイドとかBFSで解くことで各手法の理解を深めるという記事です。

　前回のABC160で出されたこのD問題

　要は、

　こういうグラフがあったとします。赤の距離はすべて1です。この場合、1と2の距離は1、1と5の距離は4、つまりiとjの距離は|i-j|であることはほぼ自明かと思います。一筋縄でいかないのはこれに、

　このようなショートカットが加わった時です。この場合、1と5の距離は3になります。こういった状況のもと、全ての頂点の組み合わせに対して最短距離を列挙しなさい、というのが問題の要旨です。

想定解

　今回は、ショートカットが1個だけなので、ショートカットを使う場合と使わない場合で比較するだけで最短経路が出ます。
　n個の頂点からiとjを選び、それぞれに対してショートカット使用/不使用の比較をするだけなので、計算量は$O(n^2)$で済みます。
　制約条件は$N \leq 2*10^3$とグラフ問題にしては大きめですが、$O(n^2)$ならこれで間に合います。

　以下は私の実装です。

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for(int i = 0; i < n; i++)
#define rep1(i, n) for(int i = 1; i < n+1; i++)
#define sort(A) sort(A.begin(),A.end())
#define reverse(A) reverse(A.begin(),A.end());
typedef long long ll;

int main(){
  int n,x,y;
  cin >> n >> x >> y;
  vector<int> ans(n);
  rep1(i,n){//i番目の頂点と
    for(int j=i+1;j<=n;j++){//j番目の頂点
      int path1 = j-i;//ショートカットを使わない場合

      //ショートカットを使う場合
      int path2 = abs(x-i);//xに着くまで
      path2 ++;//yに移動
      path2 += abs(y-j);//jに移動

      ans[min(path1,path2)] ++;
    }
  }
  rep1(i,n-1) cout << ans[i] << endl;
}

　最大実行時間$8ms$でACできました。

ワーシャルフロイド

　グラフの最短経路を出す有名な方法としてワーシャルフロイド法というものがあります。名前はごついですが、機能はシンプルです。ある隣接行列が与えられた時に、それを最短経路で更新するものです。

　隣接行列とは、すべての頂点を縦と横に並べてその頂点間の距離を行列にしたものです。当然、対称行列になり（ijの距離とjiの距離は同じ）、また対角要素はすべて0になります（自分自身への距離は0）。

　上の画像に上げた例を隣接行列にします。実装は、

int main(){
  int n,x,y;
  cin >> n >> x >> y;
  vector<vector<int>> d(n+1,vector<int>(n+1,0));
  rep1(i,n){
    rep1(j,n){
      d[i][j] = abs(i-j);
    }
  }
  d[x][y] = 1;
  d[y][x] = 1;

　となり、この段階で隣接行列dは、

\begin{matrix}
0 & 1 & 2 & 3 & 4 \\
1 & 0 & 1 & 1 & 3 \\
2 & 1 & 0 & 1 & 2 \\
3 & 1 & 1 & 0 & 1 \\
4 & 3 & 2 & 1 & 0 
\end{matrix}

　となります。

　$(2,4)$と$(4,2)$の値は2ではなく1になっていますが、これはショートカットを示しています。しかし、ショートカットを経由すれば$(1,5)$は3になるはずなのにそうなっていません。

　そこで、

  rep1(k,n){
    rep1(i,n){
      rep1(j,n){
        a[i][j] = min(d[i][j], d[i][k] + d[k][j]);
      }
    }
  }

　なる変換（ワーシャルフロイド法）を行います。
　結果、dは以下のように更新されます。

\begin{matrix}
0 & 1 & 2 & 2 & 3 \\
1 & 0 & 1 & 1 & 2 \\
2 & 1 & 0 & 1 & 2 \\
2 & 1 & 1 & 0 & 1 \\
3 & 2 & 2 & 1 & 0 
\end{matrix}

　これで最短距離が求まりました。後はこれに出てくる答えを数え上げれば正解となります。

　以下は私の実装です。

int main(){
  int n,x,y;
  cin >> n >> x >> y;
  vector<vector<int>> d(n+1,vector<int>(n+1,0));
  rep1(i,n){
    rep1(j,n){
      d[i][j] = abs(i-j);
    }
  }
  d[x][y] = 1;
  d[y][x] = 1;

  rep1(k,n){
    rep1(i,n){
      rep1(j,n){
        d[i][j] = min(d[i][j], d[i][k] + d[k][j]);
      }
    }
  }

  vector<int> ans(n);
  rep1(i,n) rep1(j,n) ans[d[i][j]] ++;
  rep1(i,n-1) cout << ans[i]/2 << endl;//重複して数えているので÷2
}

　しかしこれを提出すると……

　TLE!　あとなんかメモリもすごいことになってる！

　ワーシャルフロイドは3重ループを回しているため、今回の制約だと時間がかかりすぎるのです。

BFS（幅優先探索）

　最短経路を求めるのに有名な手段としてもう1つ、BFS（Breadth First Search、幅優先探索）というものがあります。これは迷路を探索する場合において、すべての分かれ道に対して歩数が同じになるようにじわりじわりと進めていく、という方法です。ある地点への歩数は最短経路になります。

　頂点iと頂点jの組み合わせに対してすべてBFSを行えば最短経路は求まるはずです。

　今回は隣接行列ではなく隣接リストという形式でグラフデータを格納します。この例だと、

to[1] = 2
to[2] = 1,3,4
to[3] = 2,4
to[4] = 2,3,5
to[5] = 4

　というような配列になります。

　以下は私の実装です。

int main(){
  int n,x,y;
  cin >> n >> x >> y;

  vector<vector<int>> to(n+1,vector<int>(0));

  rep1(i,n-1){
    to[i].push_back(i+1);
    to[i+1].push_back(i);
  }
  to[x].push_back(y);
  to[y].push_back(x);

  vector<int> ans(n,0);

  rep1(i,n){
    rep1(j,n){

      //こっからBFS
      int INF = 1000000000;
      vector<int> d(n+1,INF);
      vector<bool> seen(n+1,false);
      queue<int> q;
      q.push(i);
      d[i] = 0;
      while(!q.empty()){
        int now = q.front();
        q.pop();
        if(seen[now]) continue;
        seen[now] = true;
        for(int next : to[now]){
          d[next] = min(d[now]+1,d[next]);
          q.push(next);
        }
      }//while
      ans[d[j]] ++;
    }//j
  }//i

  rep1(i,n-1) cout << ans[i]/2 << endl;//重複して数えているので÷2
}

　しかしこれも、

　TLEします。

　BFSの計算量が$O(N+M)$（※Mは辺の数）なので、2重ループでBFSを回すと3重ループになってしまうのです。

改良

　しかしよく考えてみると、iのループを1巡回した時点で、距離を格納した配列（上のコードだとd）はiを起点にした最短距離で埋め尽くされています。なので、iとjについてすべてBFSを試す必要はありません。iについてBFSを行った後で、それを元にjのループを回せば、計算は2重ループまでで済みます。

　以下は私の実装です。

int main(){
  int n,x,y;
  cin >> n >> x >> y;

  vector<vector<int>> to(n+1,vector<int>(0));

  rep1(i,n-1){
    to[i].push_back(i+1);
    to[i+1].push_back(i);
  }
  to[x].push_back(y);
  to[y].push_back(x);

  vector<int> ans(n,0);

  rep1(i,n){

    //こっからBFS
    int INF = 1000000000;
    vector<int> d(n+1,INF);
    vector<bool> seen(n+1,false);
    queue<int> q;
    q.push(i);
    d[i] = 0;
    while(!q.empty()){
      int now = q.front();
      q.pop();
      if(seen[now]) continue;
      seen[now] = true;
      for(int next : to[now]){
        d[next] = min(d[now]+1,d[next]);
        q.push(next);
      }
    }//while
    rep1(j,n) ans[d[j]] ++;
  }//i

  rep1(i,n-1) cout << ans[i]/2 << endl;//重複して数えているので÷2
}

　AC！　やったぜ。