366. Find Leaves of Binary Tree

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:

          []         

/**
 * 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>> ans;

        while (root != NULL) {
            ans.push_back(helper(root));
        }
        return ans;
    }
   
    vector<int> helper(TreeNode*& root) {
        vector<int> ans;
        if (root == NULL) {
            return ans;
        }

        queue<pair<TreeNode *, TreeNode *>> q;
        q.push({root, NULL});
        while (!q.empty()) {
            pair<TreeNode *, TreeNode *> node = q.front(); q.pop();
            if (node.first -> left == NULL && node.first -> right == NULL) {
                ans.push_back(node.first->val);
                if (node.second != NULL) {
                    if (node.second -> left == node.first) {
                        node.second -> left = NULL;
                    } else if (node.second -> right == node.first) {
                        node.second -> right = NULL;
                    }
                } else {
                    root = NULL;
                }
            } else {
                if (node.first -> left != NULL) {
                    q.push({node.first -> left, node.first});
                }
                if (node.first -> right != NULL) {
                    q.push({node.first -> right, node.first});
                }
            }
        }
        return ans;
    }
};