【Gym - 100812G 】Short Path （SPFA）

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <vector>
#include <queue>
#define inf 0x3f3f3f3f3f3f3f3f
#define N 100005
#define ll long long
using namespace std;

typedef pair<ll,int> PII;
vector<PII> edge[N<<1];

int n,m,a[N];
ll d[N];
int near[N];
bool b[N];
queue <int> q;

void SPFA(){
	for(int i=1;i<=n;++i)if(!d[i]) {
		q.push(i);
		b[i]=true;
		near[i]=i;
	}
	while(!q.empty()){
		int u=q.front(); q.pop(); b[u]=false;
		for(int j=0;j<edge[u].size();++j){
			int v = edge[u][j].second;
			ll w = edge[u][j].first;
			if(d[u]+w<d[v]){
				d[v]=d[u]+w;
				near[v]=near[u];
				if(!b[v]){
					q.push(v);
					b[v]=true;
				}
			}
		}
	}
}
int main() {
	scanf("%d %d",&n,&m);
	for(int i=1;i<=n;++i){
		scanf("%d",&a[i]);
		d[i]=a[i]?0:inf;
	}
	for(int i=1;i<=m;++i){
		int u,v,w;
		scanf("%d%d%d",&u,&v,&w);
		edge[u].push_back(make_pair(w,v));
		edge[v].push_back(make_pair(w,u));
	}
	SPFA();
	int ansa,ansb;
	ll ans=inf;
	for(int i=1;i<=n;++i)
	for(int j=0;j<edge[i].size();++j){
		int to = edge[i][j].second;
		ll w = edge[i][j].first;
		if(a[i]&&a[to]){
			if(ans>w){
				ans=w;
				ansa=i;
				ansb=to;
			}
		}else if(near[i]!=near[to]){
			if(ans>d[i]+d[to]+w){
				ans=d[i]+d[to]+w;
				ansa=near[i];
				ansb=near[to];
			}
		}
	}
	if(ans!=inf)printf("%lld\n%d %d",ans,ansa,ansb);
	else puts("No luck at all");
	return 0;
}