import operator
import math

class DecisionTree:
    def __init__(self):
        pass

    # 加載數(shù)據(jù)集
    def loadData(self):
        # 天氣晴(2),陰(1),雨(0);溫度炎熱(2),適中(1),寒冷(0);濕度高(1),正常(0)
        # 風速強(1),弱(0);進行活動(yes),不進行活動(no)
        # 創(chuàng)建數(shù)據(jù)集
        data = [
            [2, 2, 1, 0, "yes"],
            [2, 2, 1, 1, "no"],
            [1, 2, 1, 0, "yes"],
            [0, 0, 0, 0, "yes"],
            [0, 0, 0, 1, "no"],
            [1, 0, 0, 1, "yes"],
            [2, 1, 1, 0, "no"],
            [2, 0, 0, 0, "yes"],
            [0, 1, 0, 0, "yes"],
            [2, 1, 0, 1, "yes"],
            [1, 2, 0, 0, "no"],
            [0, 1, 1, 1, "no"],
        ]
        # 分類屬性
        features = ["天氣", "溫度", "濕度", "風速"]
        return data, features

    # 計算給定數(shù)據(jù)集的香農(nóng)熵
    def ShannonEnt(self, data):
        numData = len(data)  # 求長度
        labelCounts = {}
        for feature in data:
            oneLabel = feature[-1]  # 獲得標簽
            # 如果標簽不在新定義的字典里創(chuàng)建該標簽值
            labelCounts.setdefault(oneLabel, 0)
            # 該類標簽下含有數(shù)據(jù)的個數(shù)
            labelCounts[oneLabel] += 1
        shannonEnt = 0.0
        for key in labelCounts:
            # 同類標簽出現(xiàn)的概率
            prob = float(labelCounts[key]) / numData
            # 以2為底求對數(shù)
            shannonEnt -= prob * math.log2(prob)
        return shannonEnt

    # 劃分數(shù)據(jù)集,三個參數(shù)為帶劃分的數(shù)據(jù)集，劃分數(shù)據(jù)集的特征，特征的返回值
    def splitData(self, data, axis, value):
        retData = []
        for feature in data:
            if feature[axis] == value:
                # 將相同數(shù)據(jù)集特征的抽取出來
                reducedFeature = feature[:axis]
                reducedFeature.extend(feature[axis + 1 :])
                retData.append(reducedFeature)
        return retData  # 返回一個列表

    # 選擇最好的數(shù)據(jù)集劃分方式
    def chooseBestFeatureToSplit(self, data):
        numFeature = len(data[0]) - 1
        baseEntropy = self.ShannonEnt(data)
        bestInfoGain = 0.0
        bestFeature = -1
        for i in range(numFeature):
            # 獲取第i個特征所有的可能取值
            featureList = [result[i] for result in data]
            # 從列表中創(chuàng)建集合，得到不重復的所有可能取值
            uniqueFeatureList = set(featureList)
            newEntropy = 0.0
            for value in uniqueFeatureList:
                # 以i為數(shù)據(jù)集特征，value為返回值，劃分數(shù)據(jù)集
                splitDataSet = self.splitData( data, i, value )
                # 數(shù)據(jù)集特征為i的所占的比例
                prob = len(splitDataSet) / float(len(data))
                # 計算每種數(shù)據(jù)集的信息熵
                newEntropy += prob * self.ShannonEnt(splitDataSet)
            infoGain = baseEntropy - newEntropy
            # 計算最好的信息增益，增益越大說明所占決策權越大
            if infoGain > bestInfoGain:
                bestInfoGain = infoGain
                bestFeature = i
        return bestFeature

    # 遞歸構建決策樹
    def majorityCnt(self, labelsList):
        labelsCount = {}
        for vote in labelsList:
            if vote not in labelsCount.keys():
                labelsCount[vote] = 0
            labelsCount[vote] += 1
        sortedLabelsCount = sorted(
            labelsCount.iteritems(), key=operator.itemgetter(1), reverse=True
        )  # 排序，True升序
        # 返回出現(xiàn)次數(shù)最多的
        print(sortedLabelsCount)
        return sortedLabelsCount[0][0]

    # 創(chuàng)建決策樹
    def createTree(self, data, features):
        # 使用"="產(chǎn)生的新變量，實際上兩者是一樣的，避免后面del()函數(shù)對原變量值產(chǎn)生影響
        features = list(features)
        labelsList = [line[-1] for line in data]
        # 類別完全相同則停止劃分
        if labelsList.count(labelsList[0]) == len(labelsList):
            return labelsList[0]
        # 遍歷完所有特征值時返回出現(xiàn)次數(shù)最多的
        if len(data[0]) == 1:
            return self.majorityCnt(labelsList)
        # 選擇最好的數(shù)據(jù)集劃分方式
        bestFeature = self.chooseBestFeatureToSplit(data)
        bestFeatLabel = features[bestFeature]  # 得到對應的標簽值
        myTree = {bestFeatLabel: {}}
        # 清空features[bestFeat],在下一次使用時清零
        del (features[bestFeature])
        featureValues = [example[bestFeature] for example in data]
        uniqueFeatureValues = set(featureValues)
        for value in uniqueFeatureValues:
            subFeatures = features[:]
            # 遞歸調(diào)用創(chuàng)建決策樹函數(shù)
            myTree[bestFeatLabel][value] = self.createTree(
                self.splitData(data, bestFeature, value), subFeatures
            )
        return myTree

    # 預測新數(shù)據(jù)特征下是否進行活動
    def predict(self, tree, features, x):
        for key1 in tree.keys():
            secondDict = tree[key1]
            # key是根節(jié)點代表的特征，featIndex是取根節(jié)點特征在特征列表的索引，方便后面對輸入樣本逐變量判斷
            featIndex = features.index(key1)
            # 這里每一個key值對應的是根節(jié)點特征的不同取值
            for key2 in secondDict.keys():
                # 找到輸入樣本在決策樹中的由根節(jié)點往下走的路徑
                if x[featIndex] == key2:
                    # 該分支產(chǎn)生了一個內(nèi)部節(jié)點，則在決策樹中繼續(xù)同樣的操作查找路徑
                    if type(secondDict[key2]).__name__ == "dict":
                        classLabel = self.predict(secondDict[key2], features, x)
                    # 該分支產(chǎn)生是葉節(jié)點，直接取值就得到類別
                    else:
                        classLabel = secondDict[key2]
        return classLabel

if __name__ == "__main__":
    dtree = DecisionTree()
    data, features = dtree.loadData()
    myTree = dtree.createTree(data, features)
    print(myTree)
    label = dtree.predict(myTree, features, [1, 1, 1, 0])
    print("新數(shù)據(jù)[1,1,1,0]對應的是否要進行活動為:{}".format(label))