вычислить высоту дерева (произвольное дерево) — python

Учитывая массив в качестве входных данных, мы должны вычислить высоту дерева. Например, введите [4, -1, 4, 1, 1] означает, что имеется 5 узлов с номерами от 0 до 4, узел 0 является дочерним по отношению к узлу 4, узел 1 является корнем, узел 2 является дочерним по отношению к узлу 4, узел 3 является дочерним элементом узла 1, а узел 4 является дочерний узел 1.

У меня очень малоэффективный код, который повторяет ввод более одного раза:

def height(array):
        di= {array.index(-1):-1}
        while (len(di)<len(array)):
                for i, p in enumerate(array):
                        if p in di.keys():
                                di[i] = di[p] - 1               
        max_distance = -1
        for val in di.values():
                if val < max_distance:
                        max_distance = val
        max_distance = abs(max_distance)
        return max_distance

Другой метод, который я пробую, — я группирую дочерние узлы вместе следующим образом: input [9, 7, 5, 5, 2, 9, 9, 9, 2, -1] Затем я группирую в {9: [0, 5, 6, 7], 7: [1], 5: [2, 3], 2: [4, 8], -1: [9]} Но я застрял и понятия не имею, на правильном ли я пути и что делать после этого. Пожалуйста, дай мне знать, что ты думаешь. Очень признателен!

• Стандартный отступ — четыре пробела.
• Пропавшие места тут и там.
• Ненужные скобки после while.
• Не справляется с пустыми деревьями.
• имя di действительно непонятно, i и p тоже могло быть лучше.
• С помощью .keys() вроде это бессмысленно.
• Использование отрицательных значений бессмысленно и сбивает с толку. С помощью abs результат особенно сбивает с толку, поскольку предполагает, что значение, которое вы ему даете, может быть отрицательным или положительным. Поскольку вы знаете, что это отрицательно, это должно было быть просто -max_distance.
• Не нужно изобретать заново max.

Соответственно переписать:

def height(array):
    depth = {-1: 0}
    while len(depth) <= len(array):
        for child, parent in enumerate(array):
            if parent in depth:
                depth[child] = depth[parent] + 1
    return max(depth.values())

Для повышения производительности вместо многократного поиска детей с готовыми родителями сначала выполните один проход, чтобы собрать детей каждого из родителей. Затем, когда вы закончите родительский элемент, просто просмотрите его ранее собранные дочерние элементы, чтобы завершить их тоже. Тогда время будет линейным, а не квадратичным.

Или, если ограничения достаточно малы (пока вы еще не ответили на комментарий, спрашивающий об этом), вы можете использовать простой рекурсивный depth(node) функция украшена functools.cache, примените его ко всем узлам и возьмите макс. Также линейное время.