366. Find Leaves of Binary Tree

366. Find Leaves of Binary Tree (Medium)

Given the root of a binary tree, collect a tree's nodes as if you were doing this:

• Collect all the leaf nodes.
• Remove all the leaf nodes.
• Repeat until the tree is empty.

 

Example 1:

Input: root = [1,2,3,4,5]
Output: [[4,5,3],[2],[1]]
Explanation:
[[3,5,4],[2],[1]] and [[3,4,5],[2],[1]] are also considered correct answers since per each level it does not matter the order on which elements are returned.

Example 2:

Input: root = [1]
Output: [[1]]

 

Constraints:

• The number of nodes in the tree is in the range [1, 100].
• -100 <= Node.val <= 100

Companies:
LinkedIn, Amazon

Related Topics:
Tree, Depth-First Search, Binary Tree

Solution 1.

// OJ: https://leetcode.com/problems/find-leaves-of-binary-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(H)
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;
    }
};

Solution 2.

// OJ: https://leetcode.com/problems/find-leaves-of-binary-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(H)
class Solution {
    vector<vector<int>> ans;
    int dfs(TreeNode *root) {
        if (!root) return -1;
        int left = dfs(root->left), right = dfs(root->right), level = 1 + max(left, right);
        if (ans.size() <= level) ans.emplace_back();
        ans[level].push_back(root->val);
        return level;
    }
public:
    vector<vector<int>> findLeaves(TreeNode* root) {
        dfs(root);
        return ans;
    }
};