【Leetcode】1530. Number of Good Leaf Nodes Pairs

题目地址：

https://leetcode.com/problems/number-of-good-leaf-nodes-pairs/description/

给定一棵二叉树，求其路径长度小于等于给定值$d$的叶子节点对的个数。顺序不同也算一个点对。

可以用分治法，先递归左右子树，求得所有叶子离当前节点的距离（假设这些距离已经排好序），然后用双指针求得经过当前节点的距离小于等于$d$的叶子点对的数量，最后将整个子树的所有叶子离当前节点父亲的距离信息返回给父亲。代码如下：

class Solution {
 public:
  int countPairs(TreeNode *root, int dist) {
    int res = 0;
    dfs(root, dist, res);
    return res;
  }

  vector<int> dfs(TreeNode *cur, int dist, int &res) {
    if (!cur) return {};
    auto v1 = dfs(cur->left, dist, res), v2 = dfs(cur->right, dist, res);
    for (int i = 0, j = v2.size() - 1; i < v1.size(); i++) {
      while (j >= 0 && v1[i] + v2[j] > dist) j--;
      if (j < 0) break;
      res += j + 1;
    }

    vector<int> v;
    v.reserve(v1.size() + v2.size() + 1);
    // 如果是叶子的话，才需要将当前节点与其父亲的距离加进去，这个距离当然就是1
    if (!cur->left && !cur->right) v.push_back(1);
    // 由于已经假设了v1和v2排好序了，所以只需要二路归并就行
    for (int i = 0, j = 0; i < v1.size() || j < v2.size();) {
      if (j == v2.size() || i < v1.size() && v1[i] <= v2[j])
        v.push_back(v1[i++] + 1);
      else
        v.push_back(v2[j++] + 1);
    }
    return v;
  }
};

如果树平衡的话，时间复杂度$O(n\log n)$，具体时间取决于树的形态，空间$O(n)$。