Find Leaves of Binary Tree - LintCode

描述
给定一个二叉树，像这样收集树节点：收集并移除所有叶子，重复，直到树为空。

样例
Given binary tree:

    1
   / \
  2   3
 / \     
4   5    

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

思路

#ifndef C650_H
#define C650_H
#include<iostream>
#include<vector>
#include<set>
using namespace std;
class TreeNode{
public:
    int val;
    TreeNode *left, *right;
    TreeNode(int val){
        this->val = val;
        this->left = this->right = NULL;
    }
};
class Solution {
public:
    /*
    * @param root: the root of binary tree
    * @return: collect and remove all leaves
    */
    vector<vector<int>> findLeaves(TreeNode * root) {
        // write your code here
        vector<vector<int>> res;
        if (!root)
            return res;
        while (root)
        {
            vector<int> vec;
            root = helper(root, vec);
            res.push_back(vec);
        }
        return res;
    }
    //vec存放root的所有叶子节点,并将叶子节点删除
    TreeNode* helper(TreeNode *root,vector<int> &vec)
    {
        if (!root)
            return NULL;
        //当root为叶子节点,添加到vec,并返回NULL
        if (!root->left&&!root->right)
        {
            vec.push_back(root->val);
            return NULL;
        }
        //更新root的左右子树
        root->left = helper(root->left, vec);
        root->right = helper(root->right, vec);
        return root;
    }
};
#endif