1080. Insufficient Nodes in Root to Leaf Paths

sol.cpp

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
const int inf = 1e8;
class Solution {
public:
    TreeNode* sufficientSubset(TreeNode* root, int limit) {
        
        function<pair<bool, int>(TreeNode*, int)> dfs = [&] (TreeNode* node, int s) {
            if (node -> left == nullptr && node -> right == nullptr) {
                pair<int, bool> c;
                c.first = node -> val, c.second = s + node -> val < limit;
                return c;
            }
                

            int sum_left = -inf, sum_right = -inf;
            if (node -> left != nullptr) {
                const auto& t1 = dfs(node -> left, s + node -> val);
                sum_left = t1.first;
                if (node -> val == 2) 
                    cout << t1.first << ' ' << node -> left -> val<<"!!!"<<'\n';
                if (t1.second) node -> left = nullptr;
            }
            if (node -> right != nullptr) {
                const auto& t2 = dfs(node -> right, s + node -> val);
                sum_right = t2.first;
                if (t2.second) node -> right = nullptr;
            }
            return pair<int, bool> {node -> val + max(sum_left, sum_right), 
                                    s + node -> val + sum_left < limit && s + node -> val + sum_right < limit};
        };
        
        const auto& t = dfs(root, 0);
        return t.second ? nullptr: root;
    }
};