ダイクストラ法の高速化が出来ない

Question

問題

アルゴリズムと数学　演習問題集 80

方針

xをノードA_iとノードB_iで結ばれる長さC_iのエッジからなる無向グラフと考え、ノード0からノード(N - 1)までの最短距離を考える。もしそれがinfなら-1と出力し、そうでないならその距離を出力。

コード

import heapq
def dijkstra(graph, N): #ダイクストラ法 #N個の頂点からなるグラフの最短経路
    dist = [float('inf') for _ in range(N)]
    dist[0] = 0

    data = [[0, 0]] #未確定ノード[距離, ノード番号]
    while data:
        pos = heapq.heappop(data)[1]


        for node, cost in graph[pos]:
            if dist[pos] + cost < dist[node]:
                dist[node] = dist[pos] + cost
                heapq.heappush(data, [dist[node], node])
    return dist

N, M = map(int, input().split())

graph = [[] for _ in range(N)]
for _ in range(M):
    A, B, C = map(int, input().split())
    A -= 1
    B -= 1
    graph[A].append([B, C])
    graph[B].append([A, C])

ans = dijkstra(graph, N)[-1]
if ans == float('inf'):
    print(-1)
else:
    print(ans)

質問

慣れないながらheapを用いてダイクストラ法を実装してみました。テストケースを見ると一つの時間切れ(2000ms超え）を除けば全て正解でした。同じ頂点を何度もみないようにseen配列でそれを管理することで高速化を図ったり（しなかったり）したのですが、尚問題が解決しません。一体どのような箇所に注目して改善すればよろしいのでしょうか。素人質問にて恐縮ですが、お力添えいただける箇所がございましたら、ご教授のほど何卒よろしくお願い申し上げます。

Answer 1

手元でいろいろ試したところ、リストをタプルに変更しただけで、倍ぐらい速くなりました。

import heapq
def dijkstra(graph, N): #ダイクストラ法 #N個の頂点からなるグラフの最短経路
    dist = [float('inf') for _ in range(N)]
    dist[0] = 0

    data = [(0, 0)] #未確定ノード[距離, ノード番号]
    while data:
        pos = heapq.heappop(data)[1]


        for node, cost in graph[pos]:
            if dist[pos] + cost < dist[node]:
                dist[node] = dist[pos] + cost
                heapq.heappush(data, (dist[node], node))
    return dist

N, M = map(int, input().split())

graph = [[] for _ in range(N)]
for _ in range(M):
    A, B, C = map(int, input().split())
    A -= 1
    B -= 1
    graph[A].append((B, C))
    graph[B].append((A, C))

ans = dijkstra(graph, N)[-1]
if ans == float('inf'):
    print(-1)
else:
    print(ans)

テストデータは、以下のコードで生成したものを使用しました。

N=100000
graph = [(i, j, (j - i - 1) * 10 + 1) for i in range(5) for j in range(i + 1, N)]
print(N, len(graph))
for i, j, l in graph:
    print(i + 1, j + 1, l)

Comments

[kay_ventris4]

リストとタプルのパフォーマンスにおける性能差を知らず調べましたが、このような違いがあったのですね。実際ご提示頂いたコードは、Python3.8.2では時間切れを起こしたものの、Pypy3で1200ms程度で通りました。誠にありがとうございました。

Answer 2

全く同一のコードを、C++で提出したところ、2000ms超え → 約800msで通りました。
余程洗練された手法を用いない限りは、Pythonの宿命と受け取るしかないのでしょうか…
解決したとは言えず締め方も非常によろしくないのですが、とりあえず「自己」解決ということで完結させて頂きます。

#include <bits/stdc++.h>
#include <math.h>
using namespace std;
#define ll long long

/* infinity */
ll inf = pow(10, 60);

vector<ll> dijkstra(vector<vector<vector<ll>>> graph, ll N) {
    vector<ll> dist = {};
    for (ll i = 0; i < N; i++) {
        dist.push_back(inf);
    }
    dist.at(0) = 0;
    priority_queue<vector<ll>, vector<vector<ll>>, greater<vector<ll>>> data = {};
    data.push({0, 0});
    while (data.size() > 0) {
        ll pos = data.top().at(1);
        data.pop();
        for (vector<ll> el: graph.at(pos)) {
            ll node = el.at(0);
            ll cost = el.at(1);
            if (dist.at(pos) + cost < dist.at(node)) {
                dist.at(node) = dist.at(pos) + cost;
                data.push({dist.at(node), node});
            }
        }
    }
    return dist;
}

void solve() {
    ll N, M; cin >> N >> M;
    vector<vector<vector<ll>>> graph = {};
    for (ll i = 0; i < N; i++) {
        graph.push_back({});
    }
    for (ll i = 0; i < M; i++) {
        ll A, B, C; cin >> A >> B >> C;
        A--; B--;
        graph.at(A).push_back({B, C});
        graph.at(B).push_back({A, C});
    }

    vector<ll> dist = dijkstra(graph, N);

    ll ans = dist.at(N - 1);
    if (ans == inf) {
        cout << -1 << endl;
    }
    else {
        cout << ans << endl;
    }
}

int main() {
    solve();
}