Algo.Java.DFS.BST.CollectNodesByDepth

example

package org.example;


import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;


/**
 * Example how to collect Nodes by depth
 *
 * HashMap<Integer, List<TreeNodeInfo>> contains list of nodes on each depth
 */
public class Main {
    public static void main(String[] args) {
        TreeNode root = new TreeNode(1);
        root.left = new TreeNode(2);
        root.right = new TreeNode(3);
        root.left.right = new TreeNode(4);
        root.right.right = new TreeNode(5);

        HashMap<Integer, List<TreeNodeInfo>> map = new HashMap<>();
        collectByDepth(root, null, 0, map);
    }

    private static void collectByDepth(TreeNode node, TreeNode parent, int depth, HashMap<Integer, List<TreeNodeInfo>> result) {
        if (node == null) {
            return;
        }

        result.getOrDefault(depth, new ArrayList<>()).add(new TreeNodeInfo(node, parent));

        collectByDepth(node.left, node, depth + 1, result);
        collectByDepth(node.right, node, depth + 1, result);
    }


}

record TreeNodeInfo(TreeNode node, TreeNode parent) {
}