IPST60: Collecting

#include <bits/stdc++.h>
using namespace std;

typedef pair<int, int> pii;
typedef pair<int, pii> pipii;

const int N = 5e4;

vector<pii> adj[N + 1];
int seq[N + 1];

int main(){

    int nVertex, nEdge, tr, st, ed;
    scanf("%d%d%d%d%d", &nVertex, &nEdge, &tr, &st, &ed);
    for(int i = 1; i <= nVertex; ++i){
        scanf("%d", &seq[i]);
    }
    for(int i = 1; i <= nEdge; ++i){
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        adj[u].emplace_back(v, w);
    }

    vector<vector<int>> dist(nVertex + 1, vector<int>(tr + 1, 1e9));
    vector<vector<bool>> visited(nVertex + 1, vector<bool>(tr + 1, false));
    priority_queue<pipii, vector<pipii>, greater<pipii>> pq;
    dist[st][0] = 0;
    pq.emplace(dist[st][0], pii(st, 0));
    while(!pq.empty()){
        pii curr = pq.top().second;
        int u = curr.first;
        int s = curr.second;
        pq.pop();
        if(u == ed && s == tr){
            cout << dist[ed][tr];
            return 0;
        }
        if(visited[u][s]){
            continue;
        }
        visited[u][s] = true;
        for(pii nxt : adj[u]){
            int v = nxt.first;
            int w = nxt.second;
            int newS = (seq[v] == s + 1) ? s + 1 : s;
            if(!visited[v][newS] && dist[u][s] + w < dist[v][newS]){
                dist[v][newS] = dist[u][s] + w;
                pq.emplace(dist[v][newS], pii(v, newS));
            }
        }
    }
    cout << "-1";

    return 0;
}