实现ID3决策树的属性选择模块(竞技世界笔试题)

描述：输入的二维数组中的每一行为数据样本，每行的1~N-1列为输入属性值X，每行的最后一列为实际类别y。输出为根据ID3的信息增益划分准则对应最佳划分属性的索引。此处数据全为离散值。

示例

输入：[[1,1,1,1],[1,0,0,1],[0,1,1,0],[0,1,1,0]]；输出：0

import math
class Solution:
    def choose_best_feature(self , dataSet ):
        # write code here
        '''
        1.计算当前信息熵
        2.根据不同划分属性，计算相应的条件熵，得到对应的信息增益
        3.取出信息增益最大时对应的划分属性索引
        '''
        # 计算信息熵
        def computerE(S):
            E = 0
            for i in set(S):
                p = S.count(i) / len(S)
                E += -p * math.log(p)
            return E

        # 当前信息熵
        y = [i[-1] for i in dataSet]
        E_now = computerE(y)

        # 不同属性划分对应的条件熵
        E_if = {}
        tar = -1
        for j in range(len(dataSet[0])-1):
            x = [i[j] for i in dataSet]
            # 计算索引j对应属性的条件熵
            E_if[j] = 0
            for i in set(x):
                p = x.count(i) / len(x)
                E_if[j] +=  p * computerE([k[-1] for k in dataSet if k[j] == i])
        
        # 返回使条件熵最小（即信息增益最大）的属性索引
        return min(E_if, key=E_if.get)


print(Solution().choose_best_feature([[1,2,3,4],[2,2,1,3],[3,1,4,2],[0,1,1,0]]))  # 0
print(Solution().choose_best_feature([[1,1,1,1],[1,0,0,1],[0,1,1,0],[0,1,1,0]]))  # 0

划重点：1决策树的划分原理；2.理清字典取每列元素的操作，以及最后取字典值最大时对应的索引操作。