LeetCode: Number of Good Leaf Nodes Pairs (Midium)

問題はこちら

leetcode.com

二分木とそのルートのノードが与えられる。二つの葉のノードは、その最短距離が distance 以下だったときに「良いペア」となる。この「良いペア」となる葉の数を求めなさい。

考え方

dfsする。今見ているノードに対して左側と右側の子ノードを探索し、その左右の葉に対する距離を足して、distance 以下だったときに求めるものをインクリメントする。

dfsの関数が返すのは、「左側に存在する全ての葉への距離 + 1」と、「右側に存在する全ての葉への距離 + 1」。これは、親から自分への距離を足して返している。（親のノードに関数が帰った時を考えると理解しやすい。）

解答

class Solution {
public:
    int countPairs(TreeNode *root, int distance) {
        int ans = 0;
        
        auto rec = [&](auto &&f, TreeNode *node) -> vector<int> {
            // そもそもノードが存在しない
            if (node == nullptr) return {0};
            // 葉
            if (node->left == nullptr && node->right == nullptr) return {1};
            vector<int> l = f(f, node->left);
            vector<int> r = f(f, node->right);
            for (const auto &a : l) {
                for (const auto &b : r) {
                    // 右側と左側の距離を足して、条件を満たすか
                    if (a && b && a + b <= distance) ans++;
                }
            }
            vector<int> ret;

            // 親ノードへの伝播
            for (const auto &a : l) {
                if (a && a + 1 < distance) ret.emplace_back(a + 1);
            }
            for (const auto &b : r) {
                if (b && b + 1 < distance) ret.emplace_back(b + 1);
            }
            return ret;
        };

        rec(rec, root);
        return ans;
    }
};