[leetcode] 366. Find Leaves of Binary Tree 解题报告

题目链接: https://leetcode.com/problems/find-leaves-of-binary-tree/

Given a binary tree, find all leaves and then remove those leaves. Then repeat the previous steps until the tree is empty.

Example:
Given binary tree 

          1
         / \
        2   3
       / \     
      4   5    

Returns [4, 5, 3], [2], [1].

Explanation:

1. Remove the leaves [4, 5, 3] from the tree

          1
         / 
        2          

2. Remove the leaf [2] from the tree

          1          

3. Remove the leaf [1] from the tree

          []         

Returns [4, 5, 3], [2], [1].

思路: 一种比较直接的思路是每次截取叶子结点保存起来, 然后再从根开始遍历剪去叶子结点. 但是这样重复遍历太多, 其时间复杂度最坏可以达到O(n!), 一种更聪明的做法是每层遍历返回一个其距离叶子结点的距离, 这样就可以用类似于后根遍历的思路, 当DFS回溯的时候顺便知道了当前是处于哪一层的结点, 然后直接放到那一层的结果中.

代码如下:

/**
 * 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:
    int DFS(TreeNode* root, vector<vector<int>>& ans)
    {
        if(!root) return 0;
        int left = DFS(root->left, ans), right = DFS(root->right, ans);
        int depth = max(left, right), len = ans.size();
        if(depth < len) ans[depth].push_back(root->val);
        else ans.push_back(vector<int>{root->val});
        return depth+1;
    }
    
    vector<vector<int>> findLeaves(TreeNode* root) {
        if(!root) return {};
        vector<vector<int>> result;
        DFS(root, result);
        return result;
    }
};