决策树ID3、C4.5算法python实现

闲来无事写了个决策树，相比算法竞赛题，实现这个还算比较容易，代码如下：

import copy
#copy.deepcopy
switch = 0
def D_class(D):#某列纯属性
    return list(set(D))

def transpose(matrix):#纯训练集，无字段
        new_matrix = []
        for i in range(len(matrix[0])):
            matrix1 = []
            for j in range(len(matrix)):
                matrix1.append(matrix[j][i])
            new_matrix.append(matrix1)
        return new_matrix
#-----------------------------------------------------
def info(D):#某列纯属性
    import math
    denominator = len(D)
    l = D_class(D)
    entro = 0
    for i in l:
        counter = 0
        for j in D:
            if j == i:
                counter += 1
        r = counter/denominator
        entro -= r * math.log2(r)
    return entro

def info_c(D, A):#纯训练集，无字段
    #import math
    a = D_class(A)
    denominator = len(D)
    class_s = len(a)
    divide = []
    for i in a:
        divide.append([i,[]])
    for i in range(denominator):
        for j in range(class_s):
            if A[i] == a[j]:
                divide[j][1].append(D[i])
    entro = 0
    for i in divide:
        entro += len(i[1]) * info(i[1])/denominator
    return entro

def info_gain(D, A):
    return info(D) - info_c(D, A)

def info_gain_ratio(D, A):
    return info_gain(D, A)/info(A)
#---------------------------------------
def D_A_isequal(D, A, D_attrib):#输入为D列，与A列， A列为结果，D_attrib为D列的某一类
    a = ''; b = ''; tap = 1
    dlen = len(D)
    for i in range(dlen):
        if D[i] == D_attrib:
            if tap:
                a = D[i]
                b = A[i]
                tap = 0
            elif b != A[i]:
                return 0
    return b

def divide_data(D, loc, a):#判定字段属性相关联，需要接受带字段的训练集列表
    newd = []
    temp1 = copy.deepcopy(D[0])
    temp1.pop(loc)
    newd.append(temp1)
    newd2 = []
    for i in D[1]:
        if i[loc] == a:
            temp = copy.deepcopy(i)
            temp.pop(loc)
            newd2.append(temp)
    newd.append(newd2)
    return newd

def information_gain_ratio_max(D):#必须输入原数据只含有训练集而不含有字段的整体
    D2 = transpose(D)
    attributes = len(D2) - 1
    maxradio = 0
    attribute = 0
    for i in range(attributes):
        #print(info_gain_ratio(data2[-1], data2[i]))
        if maxradio < info_gain_ratio(D2[-1], D2[i]):
            maxradio = info_gain_ratio(D2[-1], D2[i])
            attribute = i
    return attribute

def dmt(D):#生成决策树
    max_ratio = information_gain_ratio_max(D[1])
    #print(D[0][max_ratio])
    split_attribute = D[0][max_ratio]#信息增益比最大的字段位置赋值
    dmt_list = [[split_attribute + switch]]
    D2 = transpose(D[1])
    attributes = D_class(D2[max_ratio])
    infos = {}
    #print(attributes)
    for i in attributes:
        attri = D_A_isequal(D2[max_ratio], D2[-1], i)
        if attri != 0:
            infos[i] = attri
        else:
            infos[i] = dmt(divide_data(D, max_ratio, i))
    dmt_list.append(infos)
    return dmt_list

def dmt_main():#ID3主程序
    data = []
    data1 = []
    fields = input("请依次输入字段(不输入primary_key时以'-'代替，' '为分隔符)属性：\n").split()
    fields_num = len(fields)
    data.append([i for i in range(fields_num)])
    n = eval(input("请输入训练集规模："))
    for i in range(n):
        s = input()
        data1.append(s.split())
    global switch
    switch = eval(input("是否去除primary_key的影响(1 or 0): "))
    if switch:
        for i in range(n):
            data1[i].pop(0)
    data.append(data1)
    dmt_list = dmt(data)
    print("决策树用字典与列表代替，如下：")
    print(dmt_list)
    
    m = eval(input("请输入预测集规模："))
    print("请输入%d行%d列待预测的数据的数据(不需要ID)：" % (m, (len(fields) - 2)))
    print("分别为：")
    for i in range(1, fields_num - 1):
        if i > 1: print(" ", end = '')
        print(fields[i], end = '')
    print('')
    while m:
        adata = input().split()
        adata.insert(0, '-')
        temp = dmt_list
        while isinstance(temp[1][adata[temp[0][0]]], list):
            temp = temp[1][adata[temp[0][0]]]
        print(temp[1][adata[temp[0][0]]])
        m -= 1
    #print(fields, '\n', data)

if (__name__ == "__main__"):
    dmt_main()

'''
Sample Input:

ID 年龄 有工作 有自己的房子 信贷情况 类别
15
1 青年 否 否 一般 否
2 青年 否 否 好 否
3 青年 是 否 好 是
4 青年 是 是 一般 是
5 青年 否 否 一般 否
6 中年 否 否 一般 否
7 中年 否 否 好 否
8 中年 是 是 好 是
9 中年 否 是 非常好 是
10 中年 否 是 非常好 是
11 老年 否 是 非常好 是
12 老年 否 是 好 是
13 老年 是 否 好 是
14 老年 是 否 非常好 是
15 老年 否 否 一般 否
1
2
青年 否 否 非常好
老年 否 否 非常好
    
- 日志密度L 好友密度F 是否使用真实头像H 账号是否真实R
10
1 s s no no
2 s l yes yes
3 l m yes yes
4 m m yes yes
5 l m yes yes
6 m l no yes
7 m s no no
8 l m no yes
9 m s yes yes
10 s s yes no
1
1
l m no

Sample Output:

否
否

yes

'''

效果如下：