MediumLeetCode 1334All-pairs shortest paths — Floyd-Warshall

Find the City With the Smallest Number of Neighbors (LeetCode 1334)

Problem 5 of 10 in Advanced Graphs — Practice Problems

The problem

There are n cities numbered 0 to n − 1, joined by weighted, undirected roads [u, v, w]. For each city, count how many other cities you can reach with a shortest-path distance of at most distanceThreshold. Return the city with the smallest count. If several cities tie, return the one with the largest number.

Input:  n = 4, edges = [[0,1,3], [1,2,1], [1,3,4], [2,3,1]], distanceThreshold = 4
Output: 3

City 0 reaches 1 (distance 3) and 2 (distance 4). City 3 reaches 1 (distance 2) and 2 (distance 1). Both have 2 neighbours; 1 and 2 each have 3. The tie goes to the larger number, 3.

How to think about it

Start with brute force. For each city, run a search to find shortest distances to every other city. With non-negative weights that means Dijkstra from each of the n cities, roughly O(n · E log n). That is perfectly valid, but you need the distance between every pair, and there is a simpler tool for exactly that.

The key insight: “every pair” with small n means Floyd-Warshall. Keep a matrix dist[i][j]. Then, for each city k, ask: does going through k make the trip from i to j shorter? If dist[i][k] + dist[k][j] < dist[i][j], update it. After every k has been tried as a middle stop, dist holds all shortest distances. It is three nested loops and no heap. Here n ≤ 100, so n³ is a million steps.

Then it is just counting. For each city, count the j != i with dist[i][j] <= threshold. Keep the smallest count, and on a tie prefer the later city. Scanning cities in increasing order and replacing on <= does that automatically.

Which pattern is this? All-pairs shortest paths. The tell is “for every city” or “between every pair” on a graph small enough for O(V³).

Approach

  1. Fill an n × n matrix with infinity, set dist[i][i] = 0, and copy each road weight into both dist[u][v] and dist[v][u].
  2. For each intermediate city k, for each pair (i, j), relax dist[i][j] through k.
  3. For each city i, count other cities within the threshold.
  4. Track the city with the smallest count, replacing on ties so the largest index wins.
  5. Return that city.

Solution

def findTheCity(n, edges, threshold):
    INF = float('inf')
    dist = [[INF] * n for _ in range(n)]
    for i in range(n): dist[i][i] = 0
    for u, v, w in edges: dist[u][v] = dist[v][u] = w        # undirected roads
    for kk in range(n):                                       # kk must be the outer loop
        for i in range(n):
            for j in range(n):
                if dist[i][kk] + dist[kk][j] < dist[i][j]:
                    dist[i][j] = dist[i][kk] + dist[kk][j]
    best = -1; res = 0
    for i in range(n):
        cnt = sum(1 for j in range(n) if i != j and dist[i][j] <= threshold)
        if cnt <= best or best == -1: best = cnt; res = i     # '<=' lets later cities win ties
    return res
# findTheCity(4, [[0,1,3],[1,2,1],[1,3,4],[2,3,1]], 4) -> 3

Dry run

After Floyd-Warshall the distance matrix is:

from \ to 0 1 2 3
0 0 3 4 5
1 3 0 1 2
2 4 1 0 1
3 5 2 1 0

Counting within distance 4:

City Reachable best res
0 2 2 0
1 3 2 0
2 3 2 0
3 2 2 3

The answer is 3.

Complexity

  • Time: O(n³). Three nested loops over cities dominate; counting is O(n²).
  • Space: O(n²). The distance matrix.

Edge cases and pitfalls

  • Loop order matters. The intermediate city kk must be the outermost loop. Putting it inside gives wrong distances.
  • Unreachable cities stay at infinity and are never counted.
  • Tie-breaking: the problem wants the largest index. Using < instead of <= returns the smallest.
  • Directed edges by mistake: set both dist[u][v] and dist[v][u].
  • Parallel roads between the same pair would need min when loading edges; LeetCode guarantees distinct pairs.