单源最短路 : 计数2

https://www.acwing.com/problem/content/description/385/
\(最短路与次短路\)

#include <bits/stdc++.h>
using namespace std;
#define IO ios::sync_with_stdio(false);cin.tie(0); cout.tie(0)
inline int lowbit(int x) { return x & (-x); }
#define ll long long
#define pb push_back
#define PII pair<int, int>
#define x first
#define y second
#define inf 0x3f3f3f3f
const int N = 1010, M = 20010;
int n, m, S, T;
int h[N], e[M], ne[M], w[M], idx;
int dist[N][2], cnt[N][2];
bool st[N][2];

struct Ver {
    int id, type, dist;
    bool operator> (const Ver &W) const {
        return dist > W.dist;
    }
};

void add(int a, int b, int c) {
    ne[idx] = h[a], e[idx] = b, w[idx] = c, h[a] = idx++;
}

int dijkstra1() {
    memset(dist, 0x3f, sizeof dist);
    memset(st, 0, sizeof st);
    memset(cnt, 0, sizeof cnt);
    dist[S][0] = 0, cnt[S][0] = 1;
    priority_queue<Ver, vector<Ver>, greater<Ver>> heap;
    heap.push({S, 0, 0});
    while (heap.size()) {
        Ver t = heap.top();
        heap.pop();
        int ver = t.id, type = t.type, d = t.dist, count = cnt[ver][type];
        if (st[ver][type]) continue;
        st[ver][type] = true;
        for (int i = h[ver]; ~i; i = ne[i]) {
            int j = e[i];
            if (dist[j][0] > d + w[i]) {
                dist[j][1] = dist[j][0], cnt[j][1] = cnt[j][0];
                heap.push({j, 1, dist[j][1]}); 
                dist[j][0] = d + w[i], cnt[j][0] = count;
                heap.push({j, 0, dist[j][0]});
            }
            else if (dist[j][0] == d + w[i]) cnt[j][0] += count;
            else if (dist[j][1] > d + w[i]) {
                dist[j][1] = d + w[i], cnt[j][1] = count;
                heap.push({j, 1, dist[j][1]});
            }
            else if (dist[j][1] == d + w[i]) cnt[j][1] += count;
        }
    }
    int res = cnt[T][0];
    if (dist[T][0] + 1 == dist[T][1]) res += cnt[T][1];
    return res;
}


int main() {
    IO;
    int _;
    cin >> _;
    while (_--) {
        cin >> n >> m;
        memset(h, -1, sizeof h);
        idx = 0;
        while (m--) {
            int a, b, c;
            cin >> a >> b >> c;
            add(a, b, c);
        }
        cin >> S >> T;
        cout << dijkstra1() << '\n';
    }
    return 0;
}