AtCoder Regular Contest 052 C - 高橋くんと不思議な道

問題

 N個の町がM個の道で結ばれている。道はタイプA,Bの二種類が存在し、タイプAのコストは1、タイプBのコストは(今まで通ったタイプBの道の本数) + 1である。全ての町について町0からの最小コストを求めよ。

解法

 タイプBの道はあまり多く使わない方が良い。具体的には解説PDFの通り最小の使用本数+\sqrt{N}程度である。頂点をこれらのタイプBの使用本数で拡張してダイクストラ法を行えばよいわけだが、そもそもこの使用回数がそこまで大きくならないので使用本数が小さい順にpriority_queueから出てくるようにして普通にダイクストラ法をやれば通る。

反省

 1時間53分(5WA)でAC。愚直にダイクストラ法をやってみたり(TLE)、セグメント木を使って高速化(?)してみたり(MLE)、最小全域木問題に帰着したりしないのかなーと思っていろいろ考えていたが1時間経ってしまったので解説を見た。

 なにやら前処理が面倒くさそうだなと思って、同実装するのかcreep04の提出を見てみたらなんか普通のダイクストラ法をやっているようにしか見えない。しばらく考えてから、tupleの順番がb, costの順番になっているのでそれで枝刈り相当のことが実現されているのだと気づいた。

 これで本当に上手くいくのかはきっちり証明できないが、制約の関係から枝刈りっぽいダイクストラ法をやるんだろうなというのは気づきたかったところ。ダイクストラ法は応用が結構あって難しい。

コード

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

int main() {
    ll N, M;
    cin >> N >> M;

    vector<vector<vector<ll>>> edges(2, vector<vector<ll>>(N));
    for (ll i = 0; i < M; i++) {
        ll C, A, B;
        cin >> C >> A >> B;
        edges[C][A].push_back(B);
        edges[C][B].push_back(A);
    }

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

    cost[0] = 0;
    pq.push({ 0, 0, 0 });
    while (!pq.empty()) {
        Element t = pq.top();
        pq.pop();
        if (t.cost > cost[t.node]) {
            continue;
        }

        cost[t.node] = t.cost;

        for (ll next : edges[0][t.node]) {
            ll new_cost = t.cost + 1;
            if (new_cost < cost[next]) {
                pq.push({ next, new_cost, t.b_cnt });
            }
        }
        for (ll next : edges[1][t.node]) {
            ll new_cost = t.cost + t.b_cnt + 1;
            if (new_cost < cost[next]) {
                pq.push({ next, new_cost, t.b_cnt + 1 });
            }
        }
    }

    for (ll i = 0; i < N; i++) {
        cout << cost[i] << endl;
    }
}