LeetCode 3123 - 最短路径中的边

class Solution {
public:
    using DistanceInfo = std::tuple<int, int, int, int>;  // d, node, parent, edge_index
    using PII = std::pair<int, int>;
    using TII = std::tuple<int, int, int>;

    vector<bool> findAnswer(int n, vector<vector<int>>& edges) {
        int m = edges.size();

        vector<vector<TII>> adj_list(n);  // (后继节点, 距离, 边的索引)

        for (int i = 0; i < m; ++i) {
            int u = edges[i][0], v = edges[i][1], d = edges[i][2];
            adj_list[u].emplace_back(v, d, i);
            adj_list[v].emplace_back(u, d, i);
        }

        vector<int> dis(n, -1);  // dis[i]: 0到i的最短距离
        vector<vector<PII>> pa(n);  // (到i的最短路径中i的前继, 边的索引)

        auto cmp = [](const DistanceInfo &a, const DistanceInfo &b) { return std::get<0>(a) > std::get<0>(b); };
        priority_queue<DistanceInfo, vector<DistanceInfo>, decltype(cmp)> pq(cmp);
        pq.emplace(0, 0, -1, -1);

        // 基于优先队列的Dijkstra算法
        while (!pq.empty()) {
            auto [d, u, p, e] = pq.top(); pq.pop();
            if (dis[u] != -1) {
                if (dis[u] == d) {  // 如果有多个最短路径, 都把可能的前继记录下
                    pa[u].emplace_back(p, e);
                }
                continue;
            }

            dis[u] = d; pa[u].emplace_back(p, e);
            if (u == n - 1) continue;
            for (const auto & info: adj_list[u]) {
                auto [v, d_uv, edge_idx] = info;
                if (dis[v] != -1 && dis[v] < d + d_uv) continue;
                pq.push(make_tuple(d + d_uv, v, u, edge_idx));
            }
        }

        // 根据前面的中间结果, 回溯到根节点, 对最优路径经过的边进行标记
        // 使用bfs
        vector<bool> ans(m);
        if (dis[n - 1] == -1) return ans;
        queue<int> Q;
        Q.emplace(n - 1);
        while (!Q.empty()) {
            auto node = Q.front(); Q.pop();
            for (const auto &pa_info : pa[node]) {
                auto [p, e] = pa_info;
                // 如果是根节点, 或者已经标记过, 则跳过, 避免重复入队
                if (e == -1 || ans[e]) continue;
                ans[e] = true;
                Q.push(p);
            }
        }

        return ans;
    }
};