Обрезка бинарных деревьев с нулевыми элементами

Источник задачи: https://leetcode.com/problems/binary-tree-pruning

Дан корень бинарного дерева. В его вершинах значение может быть 1 или 0.

Обрезать поддеревья этого дерева, целиком состоящие из 0.

Пример 1:
Дано: [1,null,0,0,1]
Результат: [1,null,0,null,1]


Пояснение: 
Только оранжевые составляют поддерево, целиком состоящее из 0

Пример 2:
Дано: [1,0,1,0,0,0,1]
Результат: [1,null,1,null,1]
 

Пример 3:
Дано: [1,1,0,1,1,0,1,0]
Результат: [1,1,0,1,1,null,1]

Примечения:

• В дереве не более 200 вершин.
• Значения вершины принимает 1 или 0. 

Идея решения

Лучше проверить все листья (вершины без поддеревьев) дерева. Если лист содержит значение 0, то его можно обрезать. У родителя этого листа может быть другой лист. Если он тоже 0, то мы и его обрежем, и тогда сам родитель станет листом. И его мы тоже проверим. Если же тот другой лист = 1, то мы отсавим эту ветку в покое, даже если родитель = 0.

Мы не можем применить BFS, потому что мы теряем связь с родителями. Чтобы поддерживать эту связь, нужно вводить дополнительные структуры данных, что трудоемко.

Будем использовать DFS: от корня добираемся до самого левого листа. Если он 0, то применяем рассуждения выше. Если не 0, то откатываемся на шаг назад и пытаемся углубиться в соседнюю ветку.

DFS состоит из 3 шагов:

1. проверить конечность углубления (DFS = поиск в глубину);
2. если можно углубляться, то делаем шаг углубления, получаем результат;
3. на основании результата выносим вердикт.
   

Имплементация 

class Solution {
   public TreeNode pruneTree(TreeNode root) {
// шаг 1. проверяем конечность углубления
        if ( root == null ) return null; // если корень поддерева нулевой, то обрезать нечего
       
// шаг 2. делаем углубление, получаем результат
        root.left = pruneTree (root.left); // считаем левое поддерево текущего корня
        root.right = pruneTree (root.right); // считаем правое поддерево

// шаг 3. анализируем результат шага 2 и выносим вердикт
//  собираем всё воедино:  
//  если у вершины нет левого и правого поддерева, то это лист
//  если в листе значение 0, то его можно обрезать, то есть возвращаем null   
        if ( root.val == 0 && root.left == null && root.right == null ) return null;
//   если на предыдущем шаге не вернули null
//  (это лист с 1
//  или корень с 0, но с поддеревьями/листьями с 1 хотя бы некоторыми
//  или корень с 1 и с листьями/поддеревьями с 1 хотя бы некоторыми),
//  то возвращаем этот корень.      
        return root;
    }
}