Участник:Turk0v/Find the City With the Smallest Number of Neighbors at a Threshold Distance

Материал из DISCOPAL
Перейти к: навигация, поиск
class Solution:
    def findTheCity(self, n: int, edges: List[List[int]], distanceThreshold: int) -> int:
        dict_res = defaultdict(list)
        for u, v, w in edges:
            dict_res[u].append((v, w))
            dict_res[v].append((u, w))
        fin_res = (float('inf'), float('inf'))
        for i in range(n):
            dist = [float('inf')] * n
            dist[i] = 0
            pq = [(0, i)]
            while pq:
                _, tmp_node = heapq.heappop(pq)
                for next_tmp_node, weight in dict_res[tmp_node]:
                    tmp = dist[tmp_node] + weight
                    if tmp < dist[next_tmp_node]:
                        dist[next_tmp_node] = tmp
                        heapq.heappush(pq, (tmp, next_tmp_node))
            res = sum([d <= distanceThreshold for d in dist])
            if res <= fin_res[1]:
                fin_res = (i, res)
        return fin_res[0]