树形dp, 计数

[Draymonders]

Leetcode 5474. 好叶子节点对的数量

const int N = 1024 + 10;

// dp[i][j] 表示以i为跟，叶子距离i为j的数量
int dp[N][N];

class Solution {
public:
    int res;
    int tot;

    // 返回值int 表示 rt节点的编号
    int dfs(TreeNode* rt, int k) {
        int id = ++tot;
        int l = 0, r = 0;
        if (rt->left)
            l = dfs(rt->left, k);
        if (rt->right)
            r = dfs(rt->right, k);
        // printf("id: %d\n", id);
        // printf("l: %d, r: %d\n", l, r);
        // puts("====");
        if (l == 0 && r == 0) {
            dp[id][0] = 1;
            return id;
        }

        for (int i=0; i<k; i++) {
            for (int j=0; j<k; j++) {
                // i 和 j 满足以 rt为跟的 最短路径为k的条件
                if (i+j+2 <= k) {
                    // printf("i:%d j:%d\n", i, j);
                    res += dp[l][i] * dp[r][j];
                }
            }
        }
        // puts("###");
        for (int i=0; i<k; i++)
            dp[id][i+1] += (dp[l][i] + dp[r][i]);
        return id;
    }

    int countPairs(TreeNode* root, int k) {
        if (!root) return 0;
        memset(dp, 0, sizeof(dp));
        tot = 0;
        res = 0;
        dfs(root, k);
        return res;
    }
};