[数据结构与算法]【leetcode】1530. 好叶子节点对的数量

题目链接

1530. 好叶子节点对的数量

思路分析

对于好叶子节点，就是要求两个节点间的距离小于给定值，因此，我们需要存储节点的距离，那么当访问当根节点，说明此时节点距离为0

• 此时返回0，当上一层看到叶子节点时，计算叶子节点距离本层节点距离为1，然后存储起来，对于右子树也是同样的计算方式，这样就可以得到父节点左右子树孩子的深度，然后在每一层对于大于distance的节点数量进行累加，然后在返回上一层,每一层都存储这每一叶子节点的深度。
• 在一层进行判断最终就能得到满足条件的个数

代码实现

/**
 * 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:
    vector<int> dfs(TreeNode* root, int distance, int& ans)
    {
        if(root == nullptr)
            return  {};
        //因为要统计每个节点的叶子节点深度，所以需要返回一个数组
        if(root->left == nullptr && root->right == nullptr)
            return {0};
        vector<int> ret;    //存储root的的叶子结点的深度
        vector<int> left = dfs(root->left, distance, ans);
        //这里一定要采用引用，这样得到的才是深度，不然在下面计算ans时
        for(auto& i : left)
        {
            ret.push_back(++i);
        }
        vector<int> right = dfs(root->right, distance, ans);
        for(auto& i : right)
        {
            ret.push_back(++i);
        }

        for(int i = 0; i < left.size(); ++i)
        {
            for(int j = 0; j < right.size(); ++j)
            {
                if( (left[i] + right[j]) <= distance)
                {
                    ans++;
                }
            }
        }
        return ret;
    }
    int countPairs(TreeNode* root, int distance) {
        int ans = 0;
        dfs(root, distance, ans);
        return ans;
    }
    
};