Hard Graph part 2

5321. Find the City With the Smallest Number of Neighbors at a Threshold Distance

https://leetcode.com/problems/find-the-city-with-the-smallest-number-of-neighbors-at-a-threshold-distance/

Floyd Warshall algorithm can be used to solve this problem as the problem size is relatively small.

class Solution:
    def findTheCity(self, n: int, edges: List[List[int]], distanceThreshold: int) -> int:
        dis = [[float('inf')]*n for _ in range(n)]
        for a,b, c in edges:
            dis[a][b]=dis[b][a]=c
        for i in range(n):dis[i][i]=0
        for k in range(n):
            for i in range(n):
                for j in range(n):
                    dis[i][j] = min(dis[i][j], dis[i][k]+dis[k][j])
        res = {sum(d<=distanceThreshold for d in dis[i]):i for i in range(n)}
        return res[min(res)]

We can also use the naive Dijkstra algorithm to solve this problem.

import heapq
class Solution:
    def findTheCity(self, n: int, edges: List[List[int]], distanceThreshold: int) -> int:
        cost = [[float('inf')]*n for _ in range(n)]
        graph = collections.defaultdict(list)
        for a,b,c in edges:
            cost[a][b]=cost[b][a]=c
            graph[a].append(b)
            graph[b].append(a)
            
        def dijkstra(i):
            dis = [float('inf')]*n
            dis[i], pq =0, [(0, i)]
            heapq.heapify(pq)
            while pq:
                d, i = heapq.heappop(pq)
                if d> dis[i]:continue
                for j in graph[i]:
                    this_cost = d+cost[i][j]
                    if this_cost<dis[j]:
                        dis[j] = this_cost
                        heapq.heappush(pq, (this_cost, j))
            return sum(d<=distanceThreshold for d in dis)-1
        res = {dijkstra(i):i for i in range(n)}
        return res[min(res)]