C++ iterative (stack, O(n^2), "0ms")


  • 0
    M

    For me, the trick is to keep a stack of TreeNode** instead of just TreeNode*. That way a leaf node can be deleted without needing an explicit reference to it's parent.

            vector<vector<int>> findLeaves(TreeNode* root)
            {
                vector<vector<int>> res;
                if (!root)
                    return res;
    
                vector<TreeNode**> stack;
    
                // Each pass executes a dfs and deletes the leaves.
                while (root)
                {
                    vector<int> leaves;
    
                    stack.assign(1, &root);
                    while (!stack.empty())
                    {
                        auto node = stack.back();
                        stack.pop_back();
    
                        if ((*node)->left || (*node)->right)
                        {
                            // not a leaf
                            if ((*node)->left) stack.push_back(&(*node)->left);
                            if ((*node)->right) stack.push_back(&(*node)->right);
                        }
                        else {
                            // 'delete' a leaf by setting it to null
                            leaves.push_back((*node)->val);
                            *node = nullptr;
                        }
                    }
    
                    res.push_back(move(leaves));
                }
    
                return res;
            }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.