Number of Cities at each Distance

A network consisting of M cities and M-1 roads connecting them is given. Cities are labeled with distinct integers within the range[0..(M-1)]

Roads connect cities in such a way that each pair of distinct cities is connected either by a direct road or along a path consisting of direct roads. There is exactly one way to reach any city from any other city. In other words, cities and direct roads from a tree. The number of direct roads that must be traversed is called the distance between these two cities.

Write a function that, given a non-empty zero-indexed array T consisting of M integers describing a network of M cities and M-1 roads, returns an array consisting of M-1 integers, specifying the number of cities positioned at each distance 1, 2, ..., M-1.

Array T describes a network of cities as follows:

  • If T[P] = Q and P = Q, then P is the capital;

  • If T[P] = Q and P != Q, then there is a direct road between cities P and Q

For example, given the following array T consisting of ten elements:

T[0] = 9 T[1] = 1 T[2] = 4
T[3] = 9 T[4] = 0 T[5] = 4
T[6] = 8 T[7] = 9 T[8] = 0 T[9] = 1

the function should return [1,3,2,3,0,0,0,0,0], as explained above.

Assume that:

  1. M is an integer within the range [1..100,000]

  2. Each element of array T is an integer within the range [0..M-1]

  3. There is exactly one (possibly indirect) connection between any two distinct cities

Complexity:

  1. Expected worst-case time complexity is O(M)

  2. Expected worst-case space complexity is O(M), beyond input storage(not counting the storage required for input arguments)

Complexity: O(M+E) time and O(M) space (since adjacency list is created)

import collections


def number_of_cities(T):
    if not T:
        return []

    g = collections.defaultdict(list)
    capital = 0
    for _to, _from in enumerate(T):
        if _to != _from:
            g[_from].append(_to)
        else:
            capital = _to

    distances = [0] * (len(T) - 1)
    def dfs(root, dist):
        for city in g[root]:
            distances[dist] += 1
            dfs(city, dist + 1)

    dfs(capital, 0)
    return distances


print(number_of_cities([9,1,4,9,0,4,8,9,0,1]))

Last updated