Leet Code Daily 2024-07-18 - Number of Good Leaf Node Pairs

/**
 * 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) {}
 * };
 */

class Solution {
public:
    shared_ptr<vector<int>> compute(TreeNode* root, int dis, int& ans) {
        const auto ret = make_shared<vector<int>>(dis, 0);
        if (root == nullptr) {
            return ret;
        }
        if (root->left == nullptr && root->right == nullptr) {
            (*ret)[0] = 1;
            return ret;
        }
        const auto ret_left = compute(root->left, dis, ans),
                   ret_right = compute(root->right, dis, ans);
        for (int i = dis - 2, s = (*ret_right)[0]; i >= 0; i--) {
            ans += (*ret_left)[i] * s;
            s += (*ret_right)[dis - i - 1];
        }
        for (int i = 1; i < dis; i++) {
            (*ret)[i] += (*ret_left)[i - 1] + (*ret_right)[i - 1];
        }
        return ret;
    }

    int countPairs(TreeNode* root, int distance) {
        int ans = 0;
        compute(root, distance, ans);
        return ans;
    }
};