基于信息增益和基尼指数的二叉决策树

# 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))  # 预测并打印结果

输出结果：