366. Find Leaves of Binary Tree

366. Find Leaves of Binary Tree (Medium)

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:

          []         

Solution 1.

// OJ: https://leetcode.com/problems/find-leaves-of-binary-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(logN)
class Solution {
private:
    bool dfs(TreeNode *root, vector<int> &v) {
        if (!root) return true;
        if (!root->left && !root->right) {
            v.push_back(root->val);
            return true;
        }
        if (dfs(root->left, v)) root->left = NULL;
        if (dfs(root->right, v)) root->right = NULL;
        return false;
    }
    vector<int> removeLeaves(TreeNode *root) {
        vector<int> v;
        dfs(root, v);
        return v;
    }
public:
    vector<vector<int>> findLeaves(TreeNode* root) {
        if (!root) return {};
        vector<vector<int>> ans;
        while (root->left || root->right) {
            ans.push_back(removeLeaves(root));
        }
        ans.push_back({ root->val });
        return ans;
    }
};