ARC061-E すぬけ君の地下鉄旅行/Snuke's Subway Trip

問題URL:すぬけ君の地下鉄旅行/Snuke's Subway Trip
最近ダイクストラ書きすぎて、そらで書けるようになってきた。mapmap入れる書き方は結構気に入ってるけど、枝刈りがちょい面倒なのと計算量やばそう。CFとかでこれがボトルネックで落ちたら考える。

考察

このまえおぼえた拡張ダイクストラをやればいいんだろうな

実装

思うままに書いたらTLEした。頂点数の2倍ぐらいで済むと思っていたが多分O(M^{2})なのでしょう。つらい。

TLE解(クリックで展開)

#include <bits/stdc++.h>
#define rep(i,n) for(int i=0;i<n;i++)
#define all(x) (x).begin(),(x).end()
using namespace std;
const int INF=1145141919,MOD=1e9+7;
const long long LINF=8931145141919364364,LMOD=998244353;
// const int dx[]={1,0,-1,0,1,1,-1,-1},dy[]={0,-1,0,1,1,-1,-1,1};

#define ppi pair<int,pair<int,int>>
#define fi first
#define se second
// edge[from]={to,color}
vector<pair<int,int>> edge[100010];
// cost[goal][color]
map<int,map<int,int>> cost;

int main(){
    int n,m; cin>>n>>m;
    rep(i,m){
        int p,q,c; cin>>p>>q>>c;
        edge[p].push_back({q,c});
        edge[q].push_back({p,c});
    }
    // {cost,{from,color}}
    priority_queue<ppi,vector<ppi>,greater<ppi>> pq;
    pq.push({0,{1,-1}});
    while(pq.size()){
        auto pre=pq.top(); pq.pop();
        int pre_place=pre.se.fi,pre_color=pre.se.se,pre_dist=pre.fi;
        // cout<<pre_place<<" "<<pre_color<<" "<<pre_dist<<endl;
        for(auto to:edge[pre_place]){
            int to_place=to.fi,to_color=to.se;
            if(cost[to_place][to_color]==0||pre_dist+(pre_color==to_color?0:1)<cost[to_place][to_color]){
                cost[to_place][to_color]=pre_dist+(pre_color==to_color?0:1);
                // cout<<"in:"<<to_place<<" "<<to_color<<" "<<cost[to_place][to_color]<<endl;
                pq.push({cost[to_place][to_color],{to_place,to_color}});
            }
        }
    }
    int ans=INF;
    for(auto i:cost[n]){
        ans=min(ans,i.second);
    }
    cout<<(ans==INF?-1:ans)<<endl;
    return 0;
}


結局editorialを見た。すごい。参考にして枝刈りしたらWAした。ひどい。でも方針は間違ってないはず。書き直さないで敢えてひたすら刈る。盆栽?
色々やってたらなんか通ってしまった(2137ms)
多分この行が一番効果的だったと思います。

if(pre_dist>=cost[pre_place][-1]+1) continue;
#include <bits/stdc++.h>
#define rep(i,n) for(int i=0;i<n;i++)
#define all(x) (x).begin(),(x).end()
using namespace std;
const int INF=1145141919,MOD=1e9+7;
const long long LINF=8931145141919364364,LMOD=998244353;
// const int dx[]={1,0,-1,0,1,1,-1,-1},dy[]={0,-1,0,1,1,-1,-1,1};

#define ppi pair<int,pair<int,int>>
#define fi first
#define se second
// edge[from]={to,color}
vector<pair<int,int>> edge[100010];
// cost[goal][color]
map<int,map<int,int>> cost;

int main(){
    int n,m; cin>>n>>m;
    rep(i,m){
        int p,q,c; cin>>p>>q>>c;
        edge[p].push_back({q,c});
        edge[q].push_back({p,c});
    }
    // {cost,{from,color}}
    priority_queue<ppi,vector<ppi>,greater<ppi>> pq;
    pq.push({0,{1,-1}});
    while(pq.size()){
        auto pre=pq.top(); pq.pop();
        int pre_place=pre.se.fi,pre_color=pre.se.se,pre_dist=pre.fi;
        if(cost[pre_place][-1]==0||pre_dist<cost[pre_place][-1]){
            cost[pre_place][-1]=pre_dist;
            if(pre_place!=1) pq.push({cost[pre_place][-1],{pre_place,-1}});
        }
        if(pre_dist>=cost[pre_place][-1]+1) continue;
        // cout<<pre_place<<" "<<pre_color<<" "<<pre_dist<<endl;
        for(auto to:edge[pre_place]){
            int to_place=to.fi,to_color=to.se;
            if(pre_color==-1){
                if(cost[to_place][to_color]==0||pre_dist+1<cost[to_place][to_color]){
                    cost[to_place][to_color]=pre_dist+1;
                    pq.push({cost[to_place][to_color],{to_place,to_color}});
                }
            }else{
                if(pre_color==to_color){
                    if(cost[to_place][to_color]==0||pre_dist<cost[to_place][to_color]){
                        cost[to_place][to_color]=pre_dist;
                        pq.push({cost[to_place][to_color],{to_place,to_color}});
                    }
                }
            }
        }
    }
    cout<<(cost[n][-1]==0?-1:cost[n][-1])<<endl;
    return 0;
}