Find Edges in Shortest Paths

hard graph dijkstra shortest path

Problem

An undirected weighted graph has n nodes (0-indexed) and edges [u, v, w]. Return a boolean array where answer[i] is true if edge i lies on at least one shortest path between node 0 and node n-1. If 0 and n-1 are not connected, every entry is false.

Run Dijkstra from node 0 to get distFromStart, and again from node n-1 to get distFromEnd. Edge (u, v, w) is on a shortest path iff distFromStart[u] + w + distFromEnd[v] equals the total shortest distance (checking both orientations).

Inputn = 6, edges = [[0,1,4],[0,2,1],[1,3,2],[2,3,1],[3,5,3],[2,4,5],[4,5,2]]
Output[false, true, false, true, true, false, false]
Shortest 0→5 = 5 via 0-2-3-5. Edges 1, 3, 4 lie on it; the rest do not.

import heapq

def find_answer(n, edges):
    adj = [[] for _ in range(n)]
    for i, (u, v, w) in enumerate(edges):
        adj[u].append((v, w, i))
        adj[v].append((u, w, i))
    def dijkstra(src):
        dist = [float('inf')] * n
        dist[src] = 0
        pq = [(0, src)]
        while pq:
            d, u = heapq.heappop(pq)
            if d > dist[u]:
                continue
            for v, w, _ in adj[u]:
                if d + w < dist[v]:
                    dist[v] = d + w
                    heapq.heappush(pq, (d + w, v))
        return dist
    ds = dijkstra(0)
    de = dijkstra(n - 1)
    total = ds[n - 1]
    ans = [False] * len(edges)
    if total == float('inf'):
        return ans
    for i, (u, v, w) in enumerate(edges):
        if ds[u] + w + de[v] == total or ds[v] + w + de[u] == total:
            ans[i] = True
    return ans
function findAnswer(n, edges) {
  const adj = Array.from({ length: n }, () => []);
  edges.forEach(([u, v, w], i) => {
    adj[u].push([v, w, i]); adj[v].push([u, w, i]);
  });
  function dijkstra(src) {
    const dist = new Array(n).fill(Infinity);
    dist[src] = 0;
    const pq = [[0, src]];
    while (pq.length) {
      pq.sort((a, b) => a[0] - b[0]);
      const [d, u] = pq.shift();
      if (d > dist[u]) continue;
      for (const [v, w] of adj[u]) {
        if (d + w < dist[v]) { dist[v] = d + w; pq.push([d + w, v]); }
      }
    }
    return dist;
  }
  const ds = dijkstra(0), de = dijkstra(n - 1);
  const total = ds[n - 1];
  const ans = new Array(edges.length).fill(false);
  if (total === Infinity) return ans;
  edges.forEach(([u, v, w], i) => {
    if (ds[u] + w + de[v] === total || ds[v] + w + de[u] === total) ans[i] = true;
  });
  return ans;
}
class Solution {
    public boolean[] findAnswer(int n, int[][] edges) {
        List<int[]>[] adj = new List[n];
        for (int i = 0; i < n; i++) adj[i] = new ArrayList<>();
        for (int i = 0; i < edges.length; i++) {
            int u = edges[i][0], v = edges[i][1], w = edges[i][2];
            adj[u].add(new int[]{v, w}); adj[v].add(new int[]{u, w});
        }
        long[] ds = dijkstra(n, adj, 0);
        long[] de = dijkstra(n, adj, n - 1);
        long total = ds[n - 1];
        boolean[] ans = new boolean[edges.length];
        if (total == Long.MAX_VALUE) return ans;
        for (int i = 0; i < edges.length; i++) {
            int u = edges[i][0], v = edges[i][1], w = edges[i][2];
            if (ds[u] + w + de[v] == total || ds[v] + w + de[u] == total) ans[i] = true;
        }
        return ans;
    }
    long[] dijkstra(int n, List<int[]>[] adj, int src) {
        long[] dist = new long[n];
        Arrays.fill(dist, Long.MAX_VALUE);
        dist[src] = 0;
        PriorityQueue<long[]> pq = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));
        pq.offer(new long[]{0, src});
        while (!pq.isEmpty()) {
            long[] t = pq.poll();
            long d = t[0]; int u = (int) t[1];
            if (d > dist[u]) continue;
            for (int[] nb : adj[u]) {
                if (d + nb[1] < dist[nb[0]]) {
                    dist[nb[0]] = d + nb[1];
                    pq.offer(new long[]{dist[nb[0]], nb[0]});
                }
            }
        }
        return dist;
    }
}
class Solution {
public:
    vector<long long> dijkstra(int n, vector<vector<array<int,2>>>& adj, int src) {
        vector<long long> dist(n, LLONG_MAX);
        dist[src] = 0;
        priority_queue<pair<long long,int>, vector<pair<long long,int>>, greater<>> pq;
        pq.push({0, src});
        while (!pq.empty()) {
            auto [d, u] = pq.top(); pq.pop();
            if (d > dist[u]) continue;
            for (auto& nb : adj[u]) {
                if (d + nb[1] < dist[nb[0]]) {
                    dist[nb[0]] = d + nb[1];
                    pq.push({dist[nb[0]], nb[0]});
                }
            }
        }
        return dist;
    }
    vector<bool> findAnswer(int n, vector<vector<int>>& edges) {
        vector<vector<array<int,2>>> adj(n);
        for (auto& e : edges) {
            adj[e[0]].push_back({e[1], e[2]});
            adj[e[1]].push_back({e[0], e[2]});
        }
        auto ds = dijkstra(n, adj, 0), de = dijkstra(n, adj, n - 1);
        long long total = ds[n - 1];
        vector<bool> ans(edges.size(), false);
        if (total == LLONG_MAX) return ans;
        for (int i = 0; i < (int)edges.size(); i++) {
            int u = edges[i][0], v = edges[i][1], w = edges[i][2];
            if (ds[u] + w + de[v] == total || ds[v] + w + de[u] == total) ans[i] = true;
        }
        return ans;
    }
};
Time: O((V + E) log V) Space: O(V + E)