[Python]实现简单决策树

时间:2021-04-03 16:34:21

基本思路:

  通过香农熵来决定每一层使用哪一种标签做分类,分类后,通过多数表决法来决定该层两个节点的类别。每次消耗一个标签,所以一共需要递归“标签个数”层。

 # -*- coding:utf-8 -*-
import math
import operator
from collections import Counter def shannon_ent(dat):
siz = len(dat)
return 0.0 - reduce(lambda x, y: x + y,
map(lambda each: float(each)/siz * math.log(float(each)/siz, 2),
Counter(map(lambda each: each[-1], dat)).values())) def split_dataset(dat, axis, val):
ret = filter(lambda each: each[axis] == val, dat)
return map(lambda each: each[:axis]+each[axis+1:], ret) def choose_best_feature(dat):
feature_num = len(dat[0]) - 1
base_ent = shannon_ent(dat)
best_info_gain = 0.0
best_feature = -1
for i in range(feature_num):
feature_list = set([each[i] for each in dat])
cur_ent = reduce(lambda x, y: x + y,
map(lambda val: len(split_dataset(dat, i, val))/float(len(dat))*shannon_ent(split_dataset(dat, i, val)),
feature_list))
info_gain = base_ent - cur_ent
if info_gain > best_info_gain:
best_info_gain, best_feature = info_gain, i
return best_feature def majority_count(class_list):
class_dict = sorted(dict(Counter(class_list)).iteritems(), key=operator.itemgetter(1))
return class_dict[-1][0] def create_tree(dat, label):
class_list = map(lambda each: each[-1], dat)
if class_list.count(class_list[0]) == len(class_list):
return class_list[0]
if len(dat[0]) == 1:
return majority_count(class_list)
best_feature = choose_best_feature(dat)
best_label = label[best_feature]
d_tree = {best_label:{}}
del(label[best_feature])
feature_val = map(lambda each: each[best_feature], dat)
val_set = set(feature_val)
def _update_tree(val):
sub_label = label[:]
d_tree[best_label][val] = create_tree(split_dataset(dat, best_feature, val), sub_label)
map(_update_tree, val_set)
return d_tree d = [[1,1,'y'], [1,1,'y'], [1,0,'n'], [0,1,'n'], [0,1,'n']]
l = ['no surfacing', 'flippers'] print create_tree(d, l)