AtCoder Regular Contest 061E - すぬけ君の地下鉄旅行 / Snuke's Subway Trip

問題

　 $N(\le 10^5)$ 個の駅が $M(\le 2\times 10^5)$ 個の路線で繋がれている。各路線は一つの会社によって運営されており、同じ会社の路線を使い続ける限り料金は1である一方、違う会社の路線に乗り換える場合は新たに料金が1かかる。駅 $1$ から駅 $N$ まで移動するのにかかる料金の最小値を求めよ。

解法

　頂点に「最後に使った路線の会社」の情報も加えてダイクストラ法を行う。このとき普通に行うと頂点数が $O(M)$ 、一つの頂点から出ている辺の数が $O(M)$ となり、全体で辺の数が $O(M^2)$ となって間に合わない。

　各頂点について違う会社に乗り換えるときはある特殊な部分を通過するということにする。これによって辺の総数は $O(N + M)$ となり解けた。

反省

　2時間19分(8WA)でAC。最初は会社ごとにノードを作ってそれでダイクストラ法をやれば解けると勘違いしてWA(TLE, MLE含む)。これは同じ会社が管理しているものでも飛び石になったりしている場合があるし、計算量も全然十分ではなくてんでダメ。

　次はダイクストラ法の各更新部分でdfsして同じ会社の路線を埋めていくという方針でやったがWA。最短で行ける場合ではなくとも、その後に同じ会社の路線が多く続いていて有利になるパターンというのはあるので最短だけ取ってもダメ。

　この二つをやって1時間20分くらい経ってしまったので諦めて解説PDFを見たが、辺の数が $O(M^2)$ になるという説明が理解できず(辺の総数と、ある頂点における辺の数とを混同していた！)愚直な辺の構築をやってTLEというのを2回ほど。

　テストケースが公開されているのを思い出してTLEになった部分をやってみると、priority_queueのサイズが $O(M)$ になって各頂点で $O(M)$ の辺の数を調べていくのでそれはそうだということにようやく気付いた。

　あとは解説で書かれているように、乗り換えるパターンを書いてAC。あまりきれいなコードにはならなかった。

コード

#include"bits/stdc++.h"
using namespace std;
using ll = int64_t;

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);

    ll N, M;
    cin >> N >> M;

    struct Edge {
        ll to, c;
    };
    //edges[i][j] := 頂点iに路線jでいるときコスト0で行ける次の頂点
    vector<map<ll, vector<ll>>> edges(N);
    for (ll i = 0; i < M; i++) {
        ll p, q, c;
        cin >> p >> q >> c;
        p--; q--; c--;
        edges[p][c].push_back(q);
        edges[q][c].push_back(p);
    }

    constexpr ll INF = INT_MAX;
    vector<map<ll, pair<bool, ll>>> cost(N);
    struct Element {
        ll node, cost, c;
        bool operator>(const Element& rhs) const {
            return cost > rhs.cost;
        }
    };
    priority_queue<Element, vector<Element>, greater<Element>> pq;

    cost[0][-1] = { true, 0 };
    pq.push({ 0, 0, -1 });

    while (!pq.empty()) {
        Element t = pq.top();
        pq.pop();

        if (cost[t.node][t.c].first && t.cost > cost[t.node][t.c].second) {
            continue;
        }

        if (t.c == -1) {
            for (auto p : edges[t.node]) {
                ll new_cost = t.cost + 1;
                for (auto next : p.second) {
                    if (!cost[next][p.first].first || new_cost < cost[next][p.first].second) {
                        cost[next][p.first].first = true;
                        cost[next][p.first].second = new_cost;
                        pq.push({ next, new_cost, p.first });
                    }
                }
            }
        } else {
            ll new_cost = t.cost;
            for (auto next : edges[t.node][t.c]) {
                if (!cost[next][t.c].first || new_cost < cost[next][t.c].second) {
                    cost[next][t.c].first = true;
                    cost[next][t.c].second = new_cost;
                    pq.push({ next, new_cost, t.c });
                }
            }
            //-1への遷移
            if (!cost[t.node][-1].first || new_cost < cost[t.node][-1].second) {
                cost[t.node][-1].first = true;
                cost[t.node][-1].second = new_cost;
                pq.push({ t.node, new_cost, -1 });
            }

        }
    }

    ll ans = (cost[N - 1][-1].first ? cost[N - 1][-1].second : INF);

    cout << (ans == INF ? -1 : ans) << endl;
}