Задача: 863. All Nodes Distance K in Binary Tree
Сложность: medium

Дан корень бинарного дерева, значение целевого узла target и целое число k. Верните массив значений всех узлов, которые находятся на расстоянии k от целевого узла.

Ответ можно вернуть в любом порядке.

Пример:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
Output: [7,4,1]
Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

👨‍💻 Алгоритм:

1⃣Определите рекурсивную функцию add_parent(cur, parent), чтобы рекурсивно добавлять указатель на родителя к узлу cur. Если cur не пустой, добавьте указатель на parent: cur.parent = parent. Затем рекурсивно вызовите add_parent для левого и правого детей cur: add_parent(cur.left, cur) и add_parent(cur.right, cur). Вызовите add_parent(root, None), чтобы добавить все указатели на родителей (корневой узел не имеет родителя).

2⃣Инициализируйте пустой массив answer и пустое множество visited. Определите рекурсивную функцию dfs(cur, distance) для поиска всех узлов на расстоянии k от узла target. Если cur пустой или уже был посещён, вернитесь. Добавьте cur в visited, чтобы его не посещали повторно. Если distance = k, добавьте cur в answer и вернитесь.

3⃣Рекурсивно вызовите dfs для детей и родителя cur. Вызовите dfs(target, 0), чтобы найти все узлы на расстоянии k. Верните answer после завершения DFS.

😎 Решение:

class Solution {
    function distanceK($root, $target, $k) {
        $this->addParent($root, null);
        
        $answer = [];
        $visited = new SplObjectStorage();
        
        $this->dfs($target, $k, $visited, $answer);
        return $answer;
    }
    
    private function addParent($cur, $parent) {
        if ($cur) {
            $cur->parent = $parent;
            $this->addParent($cur->left, $cur);
            $this->addParent($cur->right, $cur);
        }
    }
    
    private function dfs($cur, $distance, $visited, &$answer) {
        if (!$cur || $visited->contains($cur)) return;
        $visited->attach($cur);
        if ($distance == 0) {
            $answer[] = $cur->val;
            return;
        }
        $this->dfs($cur->parent, $distance - 1, $visited, $answer);
        $this->dfs($cur->left, $distance - 1, $visited, $answer);
        $this->dfs($cur->right, $distance - 1, $visited, $answer);
    }
}

Ставь 👍 и забирай 📚 Базу знаний