如何计算二叉树中高度为h的子树总数量

首先约定规则：

• 子树定义：以任意节点为根，包含其所有后代节点构成的完整结构
• 本实现默认叶子节点高度为0，若你的场景中叶子节点高度定义为1，仅需将空节点的返回值调整为0即可，其余逻辑不变

核心思路 #

• 采用后序遍历处理二叉树：需要先获取当前节点左右子树的高度，才能计算以当前节点为根的子树高度
• 递归辅助函数返回「以当前节点为根的子树的高度」，遍历过程中同步统计高度等于目标值k的子树数量

边界处理 #

• 空节点高度默认返回-1，适配叶子节点高度为0的约定
• 输入k为负数、根节点为空时直接返回0，无需遍历

代码实现 #

首先是通用的二叉树节点定义：

typedef struct node {
    int val;
    struct node *left;
    struct node *right;
} node;

版本1：指针传递计数（无全局变量，更安全） #

// 辅助dfs函数：返回当前子树高度，同步计数符合要求的子树
int dfs(node *root, int k, int *count) {
    if (root == NULL) {
        return -1;
    }
    // 先递归计算左右子树高度
    int left_h = dfs(root->left, k, count);
    int right_h = dfs(root->right, k, count);
    // 当前子树高度 = 左右子树高度的最大值 + 1
    int cur_h = (left_h > right_h ? left_h : right_h) + 1;
    // 匹配目标高度则计数+1
    if (cur_h == k) {
        (*count)++;
    }
    return cur_h;
}

int subtree(node *root, int k) {
    int count = 0;
    // 非法输入直接返回
    if (k < 0 || root == NULL) {
        return 0;
    }
    dfs(root, k, &count);
    return count;
}

版本2：全局变量计数（写法更简洁） #

int count = 0;

int dfs(node *root, int k) {
    if (root == NULL) {
        return -1;
    }
    int left_h = dfs(root->left, k);
    int right_h = dfs(root->right, k);
    int cur_h = (left_h > right_h ? left_h : right_h) + 1;
    if (cur_h == k) {
        count++;
    }
    return cur_h;
}

int subtree(node *root, int k) {
    count = 0; // 每次调用前重置计数
    if (k < 0 || root == NULL) {
        return 0;
    }
    dfs(root, k);
    return count;
}

复杂度说明 #

• 时间复杂度：O(n)，n为二叉树节点总数，每个节点仅被遍历一次
• 空间复杂度：O(h)，h为二叉树的高度，为递归栈的空间开销；最坏情况二叉树为链状，空间复杂度为O(n)

内容的提问来源于stack exchange，提问作者Abhay Aggrawal