决策树——好瓜

import math

# 属性的取值范围(0~n-1)
data = [[0, 0, 0, 0],
        [0, 0, 0, 1],
        [0, 0, 1, 0],
        [1, 0, 1, 1],
        [0, 1, 0, 0],
        [0, 1, 0, 1],
        [1, 1, 1, 0],
        [0, 1, 1, 1],
        [0, 1, 0, 1]]
label = [1, 1, 0, 1, 0, 1, 0, 1, 1]
TARGET_LABEL = 1

ATTRIBUTION_LIMIT = [2, 2, 2, 2]
ATTRIBUTION_SIZE = len(ATTRIBUTION_LIMIT)
LABEL_LIMIT = max(label) + 1

cnt = []
for i in range(ATTRIBUTION_SIZE):
    att = []
    for j in range(ATTRIBUTION_LIMIT[i]):
        att.append([0 for _ in range(LABEL_LIMIT)])
    cnt.append(att)

for data_id in range(len(data)):
    for att_id in range(len(data[data_id])):
        t = cnt[att_id][data[data_id][att_id]][label[data_id]] + 1
        cnt[att_id][data[data_id][att_id]][label[data_id]] = t


# cnt[属性][属性的具体分类][标签]

def cal(att_id, label_id):
    res = 0.0
    att = cnt[att_id]
    tot = 0
    for comp_id in range(ATTRIBUTION_LIMIT[att_id]):
        tot = tot + sum(att[comp_id])
    for comp_id in range(ATTRIBUTION_LIMIT[att_id]):
        t = 1.0 * att[comp_id][label_id] / tot
        if t == 0.0:
            continue
        res = res + t * math.log2(t)
    return res


order = [[i, cal(i, TARGET_LABEL)] for i in range(ATTRIBUTION_SIZE)]
for i in range(ATTRIBUTION_SIZE):
    for j in range(i + 1, ATTRIBUTION_SIZE):
        if order[j][1] > order[i][1]:
            t = order[i][1]
            order[i][1] = order[j][1]
            order[j][1] = t

trie = [[-1 for _ in range( ATTRIBUTION_LIMIT[order[0][0]] )]]
trie_id = 0

def insert(data_id):
    global trie_id
    Data = data[data_id]
    p = 0
    for i in range(ATTRIBUTION_SIZE):
        att_id = order[i][0]
        t = Data[att_id]
        if trie[p][t] == -1:
            trie.append([-1 for _ in range( ATTRIBUTION_LIMIT[att_id] )])
            trie_id = trie_id + 1
            trie[p][t] = trie_id
        p = trie[p][t]
    trie[p] = label[data_id]

def query(data):
    p = 0
    for i in range(ATTRIBUTION_SIZE):
        att_id = order[i][0]
        t = data[att_id]
        if trie[p][t] == -1:
            return -1
        p = trie[p][t]
    return trie[p]


for i in range(len(data)):
    insert(i)

for i in range(len(data)):
    print(query(data[i]))