leetcode1334

 1 from heapq import heappush, heappop
 2 class Solution:
 3     def findTheCity(self, n: int, edges: ‘List[List[int]]‘, distanceThreshold: int) -> int:
 4         dis = [[float(‘inf‘)]*n for _ in range(n)]
 5         nei = collections.defaultdict(list)
 6         for i,j,x in edges:
 7             dis[i][j] = dis[j][i] = x
 8             nei[i].append(j)
 9             nei[j].append(i)
10         for i in range(n):
11             dis[i][i] = 0
12         
13         # Dijkstra
14         visited = set()
15         for i in range(n):
16             pool = [(0,i)]                                          # pool[j] = x:  d(i,j) = x    
17             while pool:
18                 x,j = heappop(pool)                                 # x = d(i,j)
19                 if (i,j) not in visited and x <= distanceThreshold: # early stop, if distance exceeds threshold
20                     visited.add((i,j))
21                     for k in nei[j]:
22                         dis[i][k] = min(dis[i][k], x+dis[j][k])     # dis(i,k) = min(dis(i,k), dis(i,j)+dis(j,k))
23                         heappush(pool, (dis[i][k], k))
24         cities = [[i,0] for i in range(n)]
25         for i in range(n):
26             for d in dis[i]:
27                 if d <= distanceThreshold:
28                     cities[i][1] += 1
29         result = sorted(cities,key=lambda x:[x[1],-x[0]])[0][0]
30         return result