# coding: UTF-8
'''
基于信息增益和基尼指数的二叉决策树的实现。
该决策树可以用于分类问题，通过选择合适的特征来划分样本。
'''
from collections import Counter
class 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.0
for label in label_counts:
prob = label_counts[label] / total
gini -= prob ** 2
return gini
def info_entropy(samples):
'''
计算信息熵。
para samples: list, 样本列表，每个样本的最后一个元素是标签。
return: float, 信息熵。
'''
label_counts = sum_of_each_label(samples)
total = len(samples)
entropy = 0.0
for label in label_counts:
prob = label_counts[label] / total
entropy -= prob * (prob * 3.321928094887362) # 以2为底的对数
return entropy
def 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, right
def 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 = gain
bestF = f
bestFvalue = fvalue
bestlson = lson
bestrson = 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.leafLabel
else:
# 否则根据特征值选择子树
sampleValue = sample[tree.f]
branch = tree.r if sampleValue >= tree.fvalue else tree.l
return 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)) # 预测并打印结果
输出结果：
