Python实现CART决策树算法及详细注释

# 计算样本属于第1个类的概率p
def calcProbabilityEnt(dataset):
    numEntries = len(dataset)
    count = 0
    label = dataset[0][len(dataset[0]) - 1]
    for example in dataset:
        if example[-1] == label:
            count += 1
    probabilityEnt = float(count) / numEntries
    return probabilityEnt

def choose_best_feature(dataset):
    # 特征总数
    numFeatures = len(dataset[0]) - 1
    # 当只有一个特征时
    if numFeatures == 1:
        return 0
    # 初始化最佳基尼系数
    bestGini = 1
    # 初始化最优特征
    index_of_best_feature = -1
    for i in range(numFeatures):
        # 去重，每个属性值唯一
        uniqueVals = set(example[i] for example in dataset)
        # 定义特征的值的基尼系数
        Gini = {}
        for value in uniqueVals:
            sub_dataset1, sub_dataset2 = split_dataset(dataset,i,value)
            prob1 = len(sub_dataset1) / float(len(dataset))
            prob2 = len(sub_dataset2) / float(len(dataset))
            probabilityEnt1 = calcProbabilityEnt(sub_dataset1)
            probabilityEnt2 = calcProbabilityEnt(sub_dataset2)
            Gini[value] = prob1 * 2 * probabilityEnt1 * (1 - probabilityEnt1) + prob2 * 2 * probabilityEnt2 * (1 - probabilityEnt2)
            if Gini[value] < bestGini:
                bestGini = Gini[value]
                index_of_best_feature = i
                best_split_point = value
    return index_of_best_feature, best_split_point