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

Intuition

Solution

from functools import total_ordering
import heapq


@total_ordering
class WeightedEdge:
    def __init__(self, p, q, weight):
        self.p = p
        self.q = q
        self.weight = weight

    def one(self):
        return self.p

    def another(self, v):
        return self.p if v == self.q else self.q

    def __eq__(self, obj):
        return obj.weight == self.weight

    def __lt__(self, obj):
        return self.weight < obj.weight


class WeightedGraph:
    def __init__(self, size):
        self.size = size
        self.adj = [[] for _ in range(size)]

    def add_edge(self, p, q, w):
        edge = WeightedEdge(p, q, w)
        self.adj[p].append(edge)
        self.adj[q].append(edge)


class Solution:
    def findTheCity(self, n: int, edges: List[List[int]], distanceThreshold: int) -> int:
        graph = WeightedGraph(n)
        for p, q, w in edges:
            graph.add_edge(p, q, w)
        min_reachable = 0
        reachable_count = n
        for vertex in range(n):
            tmp = self._reachable(graph, vertex, distanceThreshold)
            if tmp <= reachable_count:
                reachable_count = tmp
                min_reachable = vertex
        return min_reachable

    def _reachable(self, graph, vertex, threshold):
        visited = [False] * graph.size
        heap = []
        visited[vertex] = True
        for edge in graph.adj[vertex]:
            heapq.heappush(
                    heap,
                    [edge.weight, edge.another(vertex)]
                )

        while heap:
            distance, vertex = heapq.heappop(heap)
            if visited[vertex]:
                continue
            if distance > threshold:
                continue
            visited[vertex] = True
            for edge in graph.adj[vertex]:
                another = edge.another(vertex)
                if visited[another] or distance + edge.weight > threshold:
                    continue
                heapq.heappush(
                        heap,
                        [distance + edge.weight, another]
                        )
        return visited.count(True)