Dijkstra 法の実装例

はじめに

先日の　AtCoder Begineer Contest 362 の D 問題で問われていた Dijkstra 法の実装が上手くいかなかったため、いつでも振り返られるように忘備録として残す。

実装

例1: 木グラフ

特にひねりのない基本問題であれば、★コスト計算する上で、カスタマイズの可能性あり★以下の部分を適切に処理することで対応できる場合が多いと感じた。

#include <stdio.h>
#include <stdlib.h>
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vll;

// ----------------------------------------------------------------------

int main(void){
  int n, m;
  cin >> n >> m;
  vll a(n);
  for(int i = 0; i < n; i++) cin >> a[i];

  vector<vector<pair<int, ll>>> G(n);// <隣接する頂点番号, 辺の重み>
  for(int i = 0; i < m; i++) {
	int u, v;
    ll b;
    cin >> u >> v >> b;
    u--; v--;
    G[u].push_back({v, b});
    G[v].push_back({u, b});
  }

  vll dist(n, 2e18);// long long を想定した初期化

  int sv = 0; // 起点となる頂点番号
  dist[sv] = 0; // 起点そのものからの距離は、当然0

  // 現在地を表す。<スタート地点から距離(コスト), 現在地のindex>
  using plli = pair<ll, int>;
  // using pllpii = pair<ll, pii>; ※二次元配列の場合は pair 型でもつとよい
  
  priority_queue<plli, vector<plli>, greater<plli>> que;
  que.push({0, sv});// 初回データは、起点となる頂点番号となる
  
  while(!que.empty()) {
    auto [d, v] = que.top();
    que.pop();
    if (d > dist[v]) // dist[v]: 現時点でのvまでの最速距離
      continue;
    for (ll i = 0; i < G[v].size(); i++) {
      auto [nv, cost] = G[v][i];
      // ここでコスト加算する。
      // ★コスト計算する上で、カスタマイズの可能性あり★
      ll nc = d + cost + a[nv];
      if (nc < dist[nv]) {
        dist[nv] = nc;
        que.push({dist[nv], nv});
      }
    }
  }

  for(int i = 1; i < n; i++) {
  	cout << dist[i] + a[0] << " ";
  }
}

例2: グリッドグラフ

#include <stdio.h>
#include <stdlib.h>
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<vll> vvll;

// ----------------------------------------------------------------------

int main(void){
  int h, w;
  cin >> h >> w;

  vector<string> G(h);// <隣接する頂点番号, 辺の重み>
  for(int i = 0; i < h; i++) cin >> G[i];

  map<pii, ll> dist;// long long を想定した初期化
  for(int i = 0; i < h; i++) {
    for(int j = 0; j < w; j++) {
      dist[{i, j}] = 2e18;
    }
  }

  pii sv = {0, 0}; // 起点となる頂点位置
  dist[sv] = 0; // 起点そのものからの距離は、当然0

  // 現在地を表す。<スタート地点から距離(コスト), 現在地のindex>
  using pllpii = pair<ll, pii>; //現在地を表す。<スタート地点から距離(コスト), 現在地>
  priority_queue<pllpii, vector<pllpii>, greater<pllpii>> que;
  
  que.push({0, sv});// 初回データは、起点となる頂点番号となる
  
  while(!que.empty()) {
    auto [d, v] = que.top();
    que.pop();
    if (d > dist[v]) // dist[v]: 現時点でのvまでの最速距離
      continue;

    const int dx[] = {1, 0, -1, 0};
    const int dy[] = {0, 1, 0, -1};
    for(int i = 0; i < 4; i++) {
      int nx = v.second + dx[i];
      int ny = v.first + dy[i];
      if(nx < 0 || nx >= w || ny < 0 || ny >= h) continue;

      // コストの実装は適宜変更する(ここではマスに書かれている数字をコストとしている)
      int cost = G[ny][nx] - '0';

      ll nc = d + (ll)cost;
      pii nv = {ny, nx};
      if (nc < dist[nv]) {
        dist[nv] = nc;
        que.push({dist[nv], nv});
      }
    }
  }

  for(int i = 0; i < h; i++) {
    for(int j = 0; j < w; j++) {
      cout << dist[{i, j}] << " ";
    }
    cout << endl;
  }
}

終わりに

重みづけグラフの問題が出題されると、問題文を見るだけで敬遠しがちだったが、これでもう対策はバッチリだ☆

※追記(2024/10/04)

• 最小コストのパスを探索するとき、起点となる頂点番号を sv とすることで、柔軟に対応できるようにした。