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)

The Idea: First, we need to understand the graphical representation of the data. To = i, while is From = T[i]. With this in mind we can create the following tree from the earlier example. Then, following from the capital (root), we can aggregate the distances from the capital to every other city using dfs.

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]))