package com.leetcode.trees; import java.util.ArrayList; import java.util.List; import static org.junit.jupiter.api.Assertions.assertEquals; /** * Level: Medium * Problem Link: https://leetcode.com/problems/find-leaves-of-binary-tree/ * Problem Description: * Given a binary tree, collect a tree's nodes as if you were doing this: Collect and remove all leaves, repeat * until the tree is empty. * * Example: * Input: [1,2,3,4,5] * * 1 * / \ * 2 3 * / \ * 4 5 * * Output: [[4,5,3],[2],[1]] * * Explanation: * 1. Removing the leaves [4,5,3] would result in this tree: * 1 * / * 2 * * 2. Now removing the leaf [2] would result in this tree: * 1 * * 3. Now removing the leaf [1] would result in the empty tree: * [] * * @author rampatra * @since 2019-08-01 */ public class LeavesOfBinaryTree { /** * THe idea is to perform a DFS and backtrack. While backtracking, check the height of the node and insert * the node into the list indexed by their heights. * Time Complexity: * Space Complexity: * Runtime: 1 ms. * * @param root * @return */ public static List> findLeavesOfBinaryTree(TreeNode root) { List> levels = new ArrayList<>(); findLeavesOfBinaryTree(root, levels); return levels; } private static int findLeavesOfBinaryTree(TreeNode root, List> levels) { if (root == null) return -1; int leftHeight = findLeavesOfBinaryTree(root.left, levels); int rightHeight = findLeavesOfBinaryTree(root.right, levels); int height = Math.max(leftHeight, rightHeight) + 1; if (height >= levels.size()) { levels.add(height, new ArrayList<>()); } levels.get(height).add(root.val); return height; } public static void main(String[] args) { /* BST looks like: 4 / \ 1 7 / \ \ 3 8 20 / \ 2 6 */ TreeNode root = new TreeNode(4); root.left = new TreeNode(1); root.right = new TreeNode(7); root.left.left = new TreeNode(3); root.left.right = new TreeNode(8); root.left.left.left = new TreeNode(2); root.left.left.right = new TreeNode(6); root.right.right = new TreeNode(20); assertEquals("[[2, 6, 8, 20], [3, 7], [1], [4]]", findLeavesOfBinaryTree(root).toString()); } }