cities_by_distance.py

def cities_by_distance(arr):

    result = [None for _ in arr]
    temp = [0 for _ in range(len(arr) - 1)]

    for i, _ in enumerate(arr):
        origin = i
        pos = i
        hop = 0
        while pos != arr[pos]:
            if result[pos] is None:
                # not explored yet
                pos = arr[pos]
                hop += 1
            else:
                # already explored from here
                hop += result[pos]
                break

        if hop > 0:
            result[origin] = hop
            temp[hop-1] += 1

    return temp


assert cities_by_distance([9, 1, 4, 9, 0, 4, 8, 9, 0, 1]) == [1, 3, 2, 3, 0, 0, 0, 0, 0]