from __future__ import division

 import math

 table = {'age': {'young', 'middle', 'old'}, 'income': {'high', 'middle', 'low'},

          'student': {'yes', 'no'}, 'credit': {'good', 'superior'}, 'buy computer': {'yes', 'no'}}

 attrIndex = {'age': 0, 'income': 1, 'student': 2, 'credit': 3, 'buy computer': 4}

 attrList = ['age', 'income', 'student', 'credit']

 allDataSet = [

     ['young', 'high', 'no', 'good', 'no'], ['young', 'high', 'no', 'superior', 'no'],

     ['middle', 'high', 'no', 'superior', 'yes'], ['old', 'middle', 'no', 'good', 'yes'],

     ['young', 'middle', 'no', 'good', 'no'], ['young', 'low', 'yes', 'good', 'yes'],

     ['middle', 'high', 'yes', 'good', 'yes'], ['old', 'middle', 'no', 'superior', 'no'],

     ['young', 'high', 'yes', 'good', 'yes'], ['middle', 'middle', 'no', 'good', 'no']

 ]

 #求熵

 def entropy(attr, dataSet):

     valueCount = {v: {'yes': 0, 'no': 0, 'count': 0} for v in table[attr]}

     for row in dataSet:

         vName = row[attrIndex[attr]]

         decAttrVal = row[attrIndex['buy computer']] # 'yes' or 'no'

         valueCount[vName]['count'] = valueCount[vName]['count'] + 1

         valueCount[vName][decAttrVal] = valueCount[vName][decAttrVal] + 1

     infoMap = {v: 0 for v in table[attr]}

     for v in valueCount:

         if valueCount[v]['count'] == 0:

             infoMap[v] = 0

         else:

             p1 = valueCount[v]['yes'] / valueCount[v]['count']

             p2 = valueCount[v]['no'] / valueCount[v]['count']

             infoMap[v] = - ((0 if p1 == 0 else p1 * math.log(p1, 2)) + (0 if p2 == 0 else p2 * math.log(p2, 2)))

     s = 0

     for v in valueCount:

         s = s + valueCount[v]['count']

     propMap = {v: (valueCount[v]['count'] / s) for v in valueCount}

     i = 0

     for v in valueCount:

         i = i + infoMap[v] * propMap[v]

     return i

 #定义节点的数据结构

 class Node(object):

     def __init__(self, attrName):

         if attrName != '':

             self.attr = attrName

             self.childNodes = {v:Node('') for v in table[attrName]}

 #数据筛选

 def filtrate(dataSet, condition):

     result = []

     for row in dataSet:

         if row[attrIndex[condition['attr']]] == condition['val']:

             result.append(row)

     return result

 #求最大信息熵

 def maxEntropy(dataSet, attrList):

     if len(attrList) == 1:

         return attrList[0]

     else:

         attr = attrList[0]

         maxE = entropy(attr, dataSet)

         for a in attrList:

             if maxE < entropy(a, dataSet):

                 attr = a

         return attr

 #判断构建是否结束,当所有的决策属性都相等的时候,就不用在构建决策树了

 def endBuild(dataSet):

     if len(dataSet) == 1:

         return True

     buy = dataSet[0][attrIndex['buy computer']]

     for row in dataSet:

         if buy != row[attrIndex['buy computer']]:

             return False

 #构建决策树

 def buildDecisionTree(dataSet, root, attrList):

     if len(attrList) == 0 or endBuild(dataSet):

         root.attr = 'buy computer'

         root.result = dataSet[0][attrIndex['buy computer']]

         root.childNodes = {}

         return

     attr = root.attr

     for v in root.childNodes:

         childDataSet = filtrate(dataSet, {"attr":attr, "val":v})

         if len(childDataSet) == 0:

             root.childNodes[v] = Node('buy computer')

             root.childNodes[v].result = 'no'

             root.childNodes[v].childNodes = {}

             continue

         else:

             childAttrList = [a for a in attrList]

             childAttrList.remove(attr)

             if len(childAttrList) == 0:

                 root.childNodes[v] = Node('buy computer')

                 root.childNodes[v].result = childDataSet[0][attrIndex['buy computer']]

                 root.childNodes[v].childNodes = {}

             else:

                 childAttr = maxEntropy(childDataSet, childAttrList)

                 root.childNodes[v] = Node(childAttr)

                 buildDecisionTree(childDataSet, root.childNodes[v], childAttrList)

 #预测结果

 def predict(root, row):

     if root.attr == 'buy computer':

         return root.result

     root = root.childNodes[row[attrIndex[root.attr]]]

     return predict(root, row)

 rootAttr = maxEntropy(allDataSet, attrList)

 rootNode = Node(rootAttr)

 print rootNode.attr

 buildDecisionTree(allDataSet, rootNode, attrList)

 print predict(rootNode, ['old', 'low', 'yes', 'good'])