# coding: UTF-8
'''
基于信息增益和基尼指数的二叉决策树的实现。
该决策树可以用于分类问题，通过选择合适的特征来划分样本。
'''from collections import Counterclass biTree_node:'''二叉树节点定义每个节点可以是叶子节点或内部节点。'''def __init__(self, f=-1, fvalue=None, leafLabel=None, l=None, r=None, splitInfo="gini"):'''类初始化函数para f: int, 切分的特征，用样本中的特征次序表示para fvalue: float or int, 切分特征的决策值para leafLabel: int, 叶节点的标签para l: biTree_node指针, 当前节点的左子树para r: biTree_node指针, 当前节点的右子树para splitInfo: string, 切分的标准, 可取值'infogain'和'gini', 分别表示信息增益和基尼指数。每个节点都保存了其用于划分的特征以及该特征的具体值，并且指向其左右子树。如果是叶子节点，则保存了该节点的标签。'''self.f = f # 特征索引，即样本中的特征次序self.fvalue = fvalue # 特征切分值，用于决定样本走向左子树还是右子树self.leafLabel = leafLabel # 如果是叶节点，则保存对应的类别标签self.l = l # 左子树，指向当前节点的左子节点self.r = r # 右子树，指向当前节点的右子节点self.splitInfo = splitInfo # 切分标准，用于决定使用何种方法来计算最佳特征和特征值def gini_index(samples):'''计算基尼指数。para samples: list, 样本列表，每个样本的最后一个元素是标签。return: float, 基尼指数。'''label_counts = sum_of_each_label(samples)total = len(samples)gini = 1.0for label in label_counts:prob = label_counts[label] / totalgini -= prob ** 2return ginidef info_entropy(samples):'''计算信息熵。para samples: list, 样本列表，每个样本的最后一个元素是标签。return: float, 信息熵。'''label_counts = sum_of_each_label(samples)total = len(samples)entropy = 0.0for label in label_counts:prob = label_counts[label] / totalentropy -= prob * (prob * 3.321928094887362) # 以2为底的对数return entropydef split_samples(samples, feature, value):'''根据特征和值分割样本集。para samples: list, 样本列表。para feature: int, 特征索引。para value: float or int, 特征值。return: tuple, 两个列表，分别为左子集和右子集。'''left = [sample for sample in samples if sample[feature] < value]right = [sample for sample in samples if sample[feature] >= value]return left, rightdef sum_of_each_label(samples):'''统计样本中各类别标签的分布。para samples: list, 样本列表。return: dict, 标签及其出现次数的字典。'''labels = [sample[-1] for sample in samples]return Counter(labels)def build_biTree(samples, splitInfo="gini"):'''构建二叉决策树para samples: list, 样本的列表，每样本也是一个列表，样本的最后一项为标签，其它项为特征。para splitInfo: string, 切分的标准，可取值'infogain'和'gini', 分别表示信息增益和基尼指数。return: biTree_node, 二叉决策树的根节点。该函数递归地构建决策树，每次选择一个最佳特征和其值来切分样本集，直到无法有效切分为止。'''# 如果没有样本，则返回空节点if len(samples) == 0:return biTree_node()# 检查切分标准是否合法if splitInfo != "gini" and splitInfo != "infogain":return biTree_node()bestInfo = 0.0 # 最佳信息增益或基尼指数减少量bestF = None # 最佳特征bestFvalue = None # 最佳特征的切分值bestlson = None # 左子树bestrson = None # 右子树# 计算当前集合的基尼指数或信息熵curInfo = gini_index(samples) if splitInfo == "gini" else info_entropy(samples)sumOfFeatures = len(samples[0]) - 1 # 样本中特征的个数for f in range(0, sumOfFeatures): # 遍历每个特征featureValues = [sample[f] for sample in samples] # 提取特征值for fvalue in featureValues: # 遍历当前特征的每个值lson, rson = split_samples(samples, f, fvalue) # 根据特征及其值切分样本# 计算分裂后两个集合的基尼指数或信息熵if splitInfo == "gini":info = (gini_index(lson) * len(lson) + gini_index(rson) * len(rson)) / len(samples)else:info = (info_entropy(lson) * len(lson) + info_entropy(rson) * len(rson)) / len(samples)gain = curInfo - info # 计算增益或基尼指数的减少量# 找到最佳特征及其切分值if gain > bestInfo and len(lson) > 0 and len(rson) > 0:bestInfo = gainbestF = fbestFvalue = fvaluebestlson = lsonbestrson = rson# 如果找到了最佳切分if bestInfo > 0.0:l = build_biTree(bestlson, splitInfo) # 递归构建左子树r = build_biTree(bestrson, splitInfo) # 递归构建右子树return biTree_node(f=bestF, fvalue=bestFvalue, l=l, r=r, splitInfo=splitInfo)else:# 如果没有有效切分，则生成叶节点label_counts = sum_of_each_label(samples)return biTree_node(leafLabel=max(label_counts, key=label_counts.get), splitInfo=splitInfo)def predict(sample, tree):'''对给定样本进行预测para sample: list, 需要预测的样本para tree: biTree_node, 构建好的分类树return: int, 预测样本所属的类别'''if tree.leafLabel is not None: # 如果当前节点是叶节点return tree.leafLabelelse:# 否则根据特征值选择子树sampleValue = sample[tree.f]branch = tree.r if sampleValue >= tree.fvalue else tree.lreturn predict(sample, branch) # 递归下去def print_tree(tree, level='0'):'''简单打印树的结构para tree: biTree_node, 树的根节点para level: str, 当前节点在树中的深度，0表示根，0L表示左子节点，0R表示右子节点'''if tree.leafLabel is not None: # 如果是叶节点print('*' + level + '-' + str(tree.leafLabel)) # 打印标签else:print('+' + level + '-' + str(tree.f) + '-' + str(tree.fvalue)) # 打印特征索引及切分值print_tree(tree.l, level + 'L') # 打印左子树print_tree(tree.r, level + 'R') # 打印右子树if __name__ == "__main__":# 示例数据集：某人相亲的数据blind_date = [[35, 176, 0, 20000, 0],[28, 178, 1, 10000, 1],[26, 172, 0, 25000, 0],[29, 173, 2, 20000, 1],[28, 174, 0, 15000, 1]]print("信息增益二叉树：")tree = build_biTree(blind_date, splitInfo="infogain") # 构建信息增益的二叉树print_tree(tree) # 打印树结构print('信息增益二叉树对样本进行预测的结果：')test_sample = [[24, 178, 2, 17000],[27, 176, 0, 25000],[27, 176, 0, 10000]]# 对测试样本进行预测for x in test_sample:print(predict(x, tree))print("基尼指数二叉树：")tree = build_biTree(blind_date, splitInfo="gini") # 构建基尼指数的二叉树print_tree(tree) # 打印树结构print('基尼指数二叉树对样本进行预测的结果：')# 再次对测试样本进行预测for x in test_sample:print(predict(x, tree)) # 预测并打印结果
输出结果：
