/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {

    // Map<TreeNode, Integer> v_states = new HashMap<>();
    class MyTreeNode {
        TreeNode self;
        TreeNode parent;
        Integer leaf_count;
        Integer level;
        public MyTreeNode(TreeNode parent, TreeNode self, Integer leaf_count, Integer level) {
            this.parent = parent;
            this.self = self;
            this.leaf_count = leaf_count;
            this.level = level;
        }
    }

    Map<TreeNode, MyTreeNode> node2state = new HashMap<>();
    Queue<TreeNode> q = new LinkedList<>();

    public List<List<Integer>> findLeaves(TreeNode root) {
        List<List<Integer>> result_list = new ArrayList<>();

        // 1. count leaf
        // 2. link parent
        // 3. collect leafs for the 1st batch
        traverse(null, root);

        while (true) {

            if (q.size() == 0) break;

            TreeNode node = q.poll();
            MyTreeNode mynode = node2state.get(node);

            // gather result
            if (result_list.size() == mynode.level) {
                result_list.add(new ArrayList<>());
            }

            List<Integer> list = result_list.get(mynode.level);
            list.add(node.val);

            // update parent leaf_count
            MyTreeNode mynode_parent = node2state.get(mynode.parent);
            if (mynode_parent == null) continue;

            mynode_parent.leaf_count--;
            if (mynode_parent.leaf_count == 0) {
                mynode_parent.level = mynode.level + 1;
                q.offer(mynode.parent);
            }

        }

        return result_list;
    }

    public void traverse(TreeNode parent, TreeNode node) {
        if (node == null) return;

        if (node2state.containsKey(node) == false) {
            node2state.put(node, new MyTreeNode(parent, node, 0, 0));
        }

        MyTreeNode mynode = node2state.get(node);

        if (node.left != null) {
            mynode.leaf_count++;
            traverse(node, node.left);
        }

        if (node.right != null) {
            mynode.leaf_count++;
            traverse(node, node.right);
        }

        if (mynode.leaf_count == 0) {
            q.offer(node);
        }

    }

}