秘密の本棚

気になることをなんでも書きます

ダイクストラ法をC++の構造体に実装

ダイクストラ

有向グラフにおいてあるノードから別のノードへの最短経路を求める手法にダイクストラがあります。エッジのコストが0以上であるという制約はつくものの、優先度付きキューを用いることで計算量は E+V\log Vになります( Eは辺の数、 Vは頂点の数)。競プロでは基本中の基本なのでしょうか。

C++の構造体に実装

蟻本に載っているコードは、グローバル変数にグラフの情報を保存しておいていますが、これだと複数のグラフを扱いたくなったときに変数を変えたりしなくてはならず面倒です。ここではグラフを構造体として定義し、ダイクストラ法を関数として実装してみました。

struct edge{ll to, cost;};
typedef pair<ll,ll> P;
struct graph{
  ll V;
  vector<vector<edge> > G;
  vector<ll> d;

  graph(ll n){
    init(n);
  }

  void init(ll n){
    V = n;
    G.resize(V);
    d.resize(V);
    REP(i,V){
      d[i] = INF;
    }
  }

  void add_edge(ll s, ll t, ll cost){
    edge e;
    e.to = t, e.cost = cost;
    G[s].push_back(e);
  }

  void dijkstra(ll s){
    REP(i,V){
      d[i] = INF;
    }
    d[s] = 0;
    priority_queue<P,vector<P>, greater<P> > que;
    que.push(P(0,s));
    while(!que.empty()){
      P p = que.top(); que.pop();
      ll v = p.second;
      if(d[v]<p.first) continue;
      for(auto e : G[v]){
        if(d[e.to]>d[v]+e.cost){
          d[e.to] = d[v]+e.cost;
          que.push(P(d[e.to],e.to));
        }
      }
    }
  }
};

使い方

graph g(n);

と書けば n頂点のグラフが作成され、

g.add_edge(a,b,cost);

とすれば頂点aから頂点bへのコストcostの辺が張られます。そして

g.dijkstra(s);

とすれば頂点sから各頂点への最短経路をダイクストラ法で探索できます。探索結果は

int dist = g.d[i]:

と呼び出して使うことができます。

使用例

AtCoderの過去問で使いやすいのを持ってきました。
abc035.contest.atcoder.jp
こんな感じでかなりすっきり書けます。複数グラフがあってもごちゃごちゃしないのがいいですね。

int main(){
  ll n,m,t;
  cin >> n >> m >> t;
  ll a[n];
  graph g1(n); //順方向グラフ
  graph g2(n); //逆方向グラフ
  REP(i,n){
    cin >> a[i];
  }
  REP(i,m){
    ll a,b,c;
    cin >> a >> b >> c;
    g1.add_edge(a-1,b-1,c);
    g2.add_edge(b-1,a-1,c);
  }
  g1.dijkstra(0);
  g2.dijkstra(0);
  ll ans = 0;
  REP(i,n){
    if(t<g1.d[i]+g2.d[i]) continue;
    ll money = a[i]*(t-g1.d[i]-g2.d[i]);
    ans = max(money, ans);
  }
  cout << ans << endl;
  return 0;
}