Бежим по дереву #алгоритм #память #скорость

Дерево - очень полезная структура данных. Помимо того, что про него спрашивают на собеседованиях, дерево помогает хранить иерархически упорядоченные данные. Например, дерево элементов HTML, дерево зависимостей сущностей в игре, дерево подразделений в компании. Имплементация дерева лаконична и проста:

public class Tree<T>(Tree<T>? parent, T value) {
    public readonly List<Tree<T>> Children = [];
    public readonly Tree<T>? Parent = parent;
    public readonly T Value = value;

    public Tree<T> Add(T value) {
        var child = new Tree<T>(this, value);
        Children.Add(child);
        return child;
    }
}

Используя связь по ссылке (ведь Tree это класс), мы можем быстро перемещаться по дереву. Имплементация выше предельно упрощена и лишь передаёт суть происходящего.

При работе с деревом, иногда возникает необходимость собрать ВСЕ дочерние элементы. Например, нам надо ответить на вопрос сколько людей работает во всём "Департаменте IT". Это означает, что мы должны просуммировать всех людей, входящих в родительский департамент, потом всех людей из дочерних отделов, а затем всех людей в отделах, входящих в эти отделы, вплоть до самого низа организационной иерархии.

Для решения этой задачи мы можем написать метод с использованием страшного и мощного yield:

public IEnumerable<Tree<T>> GetChildren() {
    if (Children.Count == 0) yield break;

    foreach (var child in Children)
    {
        yield return child;
        foreach (var subChild in child.GetChildren(true))
        {
            yield return subChild;
        }
    }
}

Получается красиво и лакнично. Но, увы, очень прожорливо по памяти. Напомню, что yield не бесплатен и его мощь является следствием генерации объекта перечисления, который является классом, а, значит, каждое его инстанциирование загрязняет память. "Вложенность" при подобном подходе порождает ещё бОльший взрыв потребления памяти. Например, если у родителя 7 дочерних элементов, а у каждого из дочерних ещё 6, а у каждого из них ещё 5... короче, на 13700 элементах потребление памяти составляет 1.2Мб. Это много, особенно, если дерево используется постоянно.

Впрочем, человечество заметило эту проблему и придумало не только алгоритмы обхода дерева, но и их эффективные имплементации.

[ThreadStatic]
private static Stack<Tree<T>>? _traverseBuffer;

...

public IEnumerable<Tree<T>> TraverseYield() {
    var buffer = _traverseBuffer ?? new Stack<Tree<T>>(128);
    _traverseBuffer = null;

    foreach (var child in Children)
    {
        buffer.Push(child);
    }

    while (buffer.Count > 0)
    {
        var current = buffer.Pop();
        foreach (var child in current.Children)
        {
            buffer.Push(child);
        }

        yield return current;
    }

    _traverseBuffer = buffer;
}

Побный подход не только ускоряет работу (почти в 4 раза), но и позволяет почти избавиться от аллокации (48 байт против 1.2Мб).

Текст бенчмарка в комментариях. Результаты бенчмарка я тоже помещу в комментарии, так как телеграмм сужает пост, если в нём картинка. С телефона это не заметно, а вот на компьютере - очень, в результате чего код почти невозможно читать.