Find Edges in Shortest Paths
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).
n = 6, edges = [[0,1,4],[0,2,1],[1,3,2],[2,3,1],[3,5,3],[2,4,5],[4,5,2]][false, true, false, true, true, false, false]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;
}
};
Explanation
We need to flag each edge that could be part of some shortest route from node 0 to node n-1. A single Dijkstra gives distances from one source, but to decide membership we need to know the cost of the best path through each edge.
So we run Dijkstra twice: once from the start (giving ds[x], the shortest distance from 0 to every node) and once from the end (giving de[x], the shortest distance from n-1 to every node). The overall shortest distance is total = ds[n-1].
An edge (u, v, w) lies on a shortest path exactly when you can reach u optimally, cross the edge, then finish optimally to the end: ds[u] + w + de[v] == total. Since the graph is undirected we also try the reversed orientation ds[v] + w + de[u] == total.
If total is infinite the endpoints are disconnected, and we return all-false immediately. Otherwise each edge is tested in O(1) using the two precomputed distance arrays.
Example: with shortest distance 5 along 0-2-3-5, edge (0,2,1) gives 0 + 1 + 4 = 5 ✓, edge (2,3,1) gives 1 + 1 + 3 = 5 ✓, edge (3,5,3) gives 2 + 3 + 0 = 5 ✓. The other edges exceed 5, so the result marks only those three true.