title
给你一个 n 个节点的无向带权图，节点编号为 0 到 n - 1 。图中总共有 m 条边，用二维数组 edges 表示，其中 edges[i] = [ai, bi, wi] 表示节点 ai 和 bi 之间有一条边权为 wi 的边。
对于节点 0 为出发点，节点 n - 1 为结束点的所有最短路，你需要返回一个长度为 m 的 boolean 数组 answer ，如果 edges[i] 至少 在其中一条最短路上，那么 answer[i] 为 true ，否则 answer[i] 为 false 。
请你返回数组 answer 。
注意，图可能不连通。
示例 1：
输入：n = 6, edges = [[0,1,4],[0,2,1],[1,3,2],[1,4,3],[1,5,1],[2,3,1],[3,5,3],[4,5,2]]
输出：[true,true,true,false,true,true,true,false]
解释：
以下为节点 0 出发到达节点 5 的 所有 最短路：
- 路径
0 -> 1 -> 5：边权和为4 + 1 = 5。 - 路径
0 -> 2 -> 3 -> 5：边权和为1 + 1 + 3 = 5。 - 路径
0 -> 2 -> 3 -> 1 -> 5：边权和为1 + 1 + 2 + 1 = 5。
示例 2：
输入：n = 4, edges = [[2,0,1],[0,1,1],[0,3,4],[3,2,2]]
输出：[true,false,false,true]
解释：
只有一条从节点 0 出发到达节点 3 的最短路 0 -> 2 -> 3 ，边权和为 1 + 2 = 3 。
提示：
2 <= n <= 5 * 104m == edges.length1 <= m <= min(5 * 104, n * (n - 1) / 2)0 <= ai, bi < nai != bi1 <= wi <= 105- 图中没有重边。
solution
dijkstra 求最短路径上的边。也可以使用0和n-1这两个点分别dijkstra，之后对于点u，v这个u - v的边，判断u到0以及v到n-1距离就可以得出它是不是最短路径上的边。
from collections import defaultdict, Counter
MOD = int(1e9 + 7)
INF = float('inf')
import sortedcontainers
import bisect
import heapq
class Graph:
def __init__(self):
self.path = collections.defaultdict(lambda: -1)
self.path_pre = collections.defaultdict(list)
'''
边数组转邻接表
'''
def dijkstra(self, graph, start):
'''
重要的是定义成本，可能是路径，可能是花费, 如果有限制次数，也可以将次数也算进成本里
如果是次数
graph = collections.defaultdict(list)
for u,v,w in edges:
graph[u].append((v,w))
graph[v].append((u,w))
:param graph: 存了每个点的next node
'''
import heapq
heap = [(0, start, -1)] # cost from start node,end node
dist = {}
# self.path = [-1 for _ in range(len(graph.keys()))]
while heap:
(cost, u, pre) = heapq.heappop(heap)
if u in dist:
if cost == dist[u]:
self.path_pre[u].append(pre)
continue
dist[u] = cost
self.path_pre[u].append(pre)
self.path[u] = pre
for v, c in graph[u]:
if v in dist:
continue
next = cost + c
heapq.heappush(heap, (next, v, u))
return dist
def get_dijkstra_path_for_multiple(self, j):
'''
返回一个树形结构，j这个终点到起点的所有的点位以及对应的上一个点位的list
eg: defaultdict(<class 'list'>, {5: [1, 3], 1: [0, 3], 0: [-1], 3: [2], 2: [0]})
'''
nxt = collections.defaultdict(list)
def add_to_nxt(i):
if i not in self.path_pre:
return []
nxt[i] = self.path_pre[i]
for v in self.path_pre[i]:
add_to_nxt(v)
add_to_nxt(j)
return nxt
class Solution:
def findAnswer(self, n: int, edges: List[List[int]]) -> List[bool]:
graph = collections.defaultdict(list)
for u, v, w in edges:
graph[u].append((v, w))
graph[v].append((u, w))
G = Graph()
d = G.dijkstra(graph, 0)
# print(d)
pre = G.get_dijkstra_path_for_multiple(n - 1)
# print(pre)
for k in pre.keys():
pre[k] = set(pre[k])
# print(pre)
ans = []
for i,(u,v,_) in enumerate(edges):
# print(i,u,v)
if u in pre and v in pre[u] or v in pre and u in pre[v]:
ans.append(True)
else:
ans.append(False)
return ans