如何基于疾病症状数据递归构建二叉树（recursively building a binary tree）

实现前提说明 #

你提到的数据集结构为「每行首元素为疾病名称，后续元素为对应症状」，这类场景通常构建的是症状决策二叉树：内部节点存储症状判断规则，左分支代表「存在该症状」，右分支代表「不存在该症状」，叶子节点存储最终匹配的疾病名称。

第一步：数据预处理 #

• 先将原始文本文件解析为结构化存储格式，推荐用字典存储：疾病名称: 对应症状集合，示例结构如下：

disease_symptom_map = {
    "感冒": {"咳嗽", "流鼻涕", "发热", "咽喉痛"},
    "急性肠胃炎": {"腹泻", "呕吐", "腹痛", "发热"},
    "过敏性鼻炎": {"流鼻涕", "打喷嚏", "鼻塞", "眼痒"}
}

• 提取所有数据集里出现过的症状，存入可用特征列表备用。

第二步：递归构建二叉树核心逻辑 #

递归终止条件 #

满足以下任意一个条件即停止递归，生成叶子节点：

• 当前处理的数据集内所有记录都对应同一种疾病，直接将该疾病设为叶子节点值
• 没有剩余的症状特征可用于分裂，取当前数据集内出现频次最高的疾病设为叶子节点值

递归执行流程 #

1. 遍历当前所有可用的症状特征，计算每个特征的分裂增益（常用信息增益、基尼系数，用来衡量该症状对不同疾病的区分能力）
2. 选择分裂增益最高的症状作为当前内部节点的判断规则
3. 将当前数据集拆分为两个子集：
   • 左子集：所有包含当前选中症状的疾病记录
   • 右子集：所有不包含当前选中症状的疾病记录
4. 将当前使用过的症状从可用特征列表中移除，分别对左、右子集递归执行上述流程，生成当前节点的左、右子节点

简化版实现代码示例 #

import math
from collections import Counter

# 定义二叉树节点类
class TreeNode:
    def __init__(self, feature=None, left=None, right=None, disease=None):
        self.feature = feature  # 内部节点存症状特征，叶子节点为None
        self.left = left  # 存在该症状的分支
        self.right = right  # 不存在该症状的分支
        self.disease = disease  # 叶子节点存疾病名，内部节点为None

# 计算信息熵的工具函数
def calc_entropy(dataset):
    total = len(dataset)
    label_cnt = Counter([d[0] for d in dataset])
    entropy = 0.0
    for cnt in label_cnt.values():
        prob = cnt / total
        entropy -= prob * math.log2(prob)
    return entropy

# 递归构建二叉树主函数
def build_tree(dataset, available_features):
    # 终止条件1：所有样本标签相同
    labels = [d[0] for d in dataset]
    if len(set(labels)) == 1:
        return TreeNode(disease=labels[0])
    # 终止条件2：无可用特征
    if not available_features:
        max_disease = Counter(labels).most_common(1)[0][0]
        return TreeNode(disease=max_disease)
    
    # 计算所有特征的信息增益，选最优分裂特征
    best_gain = -1
    best_feature = None
    base_entropy = calc_entropy(dataset)
    for feature in available_features:
        # 按当前特征拆分数据集
        sub_left = [d for d in dataset if feature in d[1]]
        sub_right = [d for d in dataset if feature not in d[1]]
        # 计算分裂后的信息熵
        new_entropy = len(sub_left)/len(dataset)*calc_entropy(sub_left) + len(sub_right)/len(dataset)*calc_entropy(sub_right)
        gain = base_entropy - new_entropy
        if gain > best_gain:
            best_gain = gain
            best_feature = feature
    
    # 拆分数据集，移除已使用的特征
    sub_left = [d for d in dataset if best_feature in d[1]]
    sub_right = [d for d in dataset if best_feature not in d[1]]
    new_features = [f for f in available_features if f != best_feature]
    
    # 递归构建左右子树
    left_node = build_tree(sub_left, new_features)
    right_node = build_tree(sub_right, new_features)
    return TreeNode(feature=best_feature, left=left_node, right=right_node)

# 调用示例
if __name__ == "__main__":
    # 模拟输入数据集，格式为[(疾病名, 症状集合)]
    sample_dataset = [
        ("感冒", {"咳嗽", "流鼻涕", "发热", "咽喉痛"}),
        ("急性肠胃炎", {"腹泻", "呕吐", "腹痛", "发热"}),
        ("过敏性鼻炎", {"流鼻涕", "打喷嚏", "鼻塞", "眼痒"})
    ]
    # 提取所有可用症状
    all_features = set()
    for d in sample_dataset:
        all_features.update(d[1])
    all_features = list(all_features)
    # 构建二叉树
    root = build_tree(sample_dataset, all_features)

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