366. Find Leaves of Binary Tree

• Difficulty: Medium

• Topics: Tree, Depth-first Search

• Similar Questions:

Problem:

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:

          []         

Solutions:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<vector<int>> findLeaves(TreeNode* root) {
        vector<vector<int>> ret;
        visit(root, ret);
        return ret;
    }

private:
    int visit(TreeNode* root, vector<vector<int>>& ret) {
        if (root == nullptr)    return -1;
        int leftDepth = visit(root->left, ret);
        int rightDepth = visit(root->right, ret);

        int rootDepth = 1 + max(leftDepth, rightDepth);
        if (ret.size() == rootDepth) {
            ret.push_back({});
        }
        ret[rootDepth].push_back(root->val);
        return rootDepth;
    }

};