为树形结构每个节点累加统计特定类型子节点数量

问题描述 #

我被一个树形结构的统计问题卡住了，一直搞不定。现有树形结构如下：

这个结构里有两种节点：

• category（分类）：粗体显示，需要统计两个数值：
  1. 直接子分类数量（已经实现）
  2. 该分类下所有层级的**question（问题）**节点总数（搞不定这个）
• question（问题）：节点旁带圆圈标识

当前代码（存在问题） #

if (treeNode.children && treeNode.children.length) {
            treeNode.questionCount = 0;
            treeNode.subcategoryCount = treeNode.children
                .filter((item) => item.type === NODE_TYPES.CATEGORY).length;
            for (let i = 0; i < treeNode.children.length; i++) {
                treeNode.questionCount += enrichWithNumberOfSubcategoriesAndQuestions(treeNode.children[i]);
            }

        }
        if (treeNode.type === NODE_TYPES.QUESTION) {
            return 1;
        } else if (treeNode.type === NODE_TYPES.CATEGORY) {
            return treeNode.children
                .filter((item) => item.type !== NODE_TYPES.CATEGORY).length;
        }
        return 0;

问题原因 #

当前代码的核心错误在于：分类节点返回的仅仅是直接子节点中的问题数，没有递归累加子分类里的深层问题数。虽然你在循环里已经把子节点的返回值加到了treeNode.questionCount，但最后返回的却是直接子问题数，导致上层节点无法正确累加全层级的总数。

修正后的代码 #

function enrichWithNumberOfSubcategoriesAndQuestions(treeNode) {
    // 初始化统计字段，避免未定义
    treeNode.questionCount = 0;
    treeNode.subcategoryCount = 0;

    if (treeNode.children && treeNode.children.length) {
        // 统计直接子分类数量
        treeNode.subcategoryCount = treeNode.children.filter(item => item.type === NODE_TYPES.CATEGORY).length;
        
        // 递归遍历所有子节点，累加全层级问题数
        for (const child of treeNode.children) {
            treeNode.questionCount += enrichWithNumberOfSubcategoriesAndQuestions(child);
        }
    }

    // 问题节点返回1，分类节点返回累计的全层级问题数，其他节点返回0
    if (treeNode.type === NODE_TYPES.QUESTION) {
        return 1;
    } else if (treeNode.type === NODE_TYPES.CATEGORY) {
        return treeNode.questionCount;
    }
    return 0;
}

修正说明 #

1. 提前初始化两个统计字段，防止出现未定义的情况；
2. 分类节点的返回值改为已经累加好的treeNode.questionCount，这样上层节点递归时能拿到当前分类下所有层级的问题总数；
3. 循环遍历子节点时，递归调用函数并累加返回值，自然就包含了子分类里的深层问题数。

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