test

Author: lining0806

DesicionTree/DesicionTreeTest.py

@@ -6,11 +6,9 @@ class DesicionTree():
     def __init__(self):
         pass
 
-    def _calcShannonEnt(self, dataSet): ## 计算数据集的熵
-        numEntries = len(dataSet)
+    def _calcShannonEnt(self, classList): ## 计算数据集的熵
         classCounts = {}
-        for data in dataSet:
-            currentLabel = data[-1]
+        for currentLabel in classList:
             if currentLabel not in classCounts:
                 classCounts[currentLabel] = 1
             else:
@@ -24,37 +22,39 @@ def _calcShannonEnt(self, dataSet): ## 计算数据集的熵
         '''
         shannonEnt = 0.0
         for key in classCounts:
-            prob = classCounts[key]/float(numEntries)
+            prob = classCounts[key]/float(len(classList))
             shannonEnt -= prob*math.log(prob, 2) # log base 2
         return shannonEnt
 
-    def _splitDataSet(self, dataSet, axis, value):
-        retDataSet = []
-        for data in dataSet:
+    def _splitDataSet(self, dataArr, classList, axis, value):
+        retFeatData = []
+        retLabelData = []
+        for data, label in zip(dataArr, classList):
             # print data[axis]
             if data[axis] == value:
-                reduceddata = data[:axis]
-                reduceddata.extend(data[axis+1:])
-                retDataSet.append(reduceddata)
-        return retDataSet
+                reducedFeat = data[:axis]
+                reducedFeat.extend(data[axis+1:])
+                retFeatData.append(reducedFeat)
+                retLabelData.append(label)
+        return retFeatData, retLabelData
 
-    def _chooseBestFeatureToSplit(self, dataSet):
-        numFeatures = len(dataSet[0])-1 # 最后一列是类标签
-        baseEntropy = self._calcShannonEnt(dataSet)
+    def _chooseBestFeatureToSplit(self, dataArr, classList):
+        baseEntropy = self._calcShannonEnt(classList)
         bestInfoGain = 0
         bestFeature = -1
+        numFeatures = len(dataArr[0])
         for i in range(numFeatures): # 依次迭代所有的特征
-            featList = [data[i] for data in dataSet]
+            featList = [data[i] for data in dataArr]
             values = set(featList)
             '''
             条件熵：sigma(pj*子数据集的熵)
             '''
             ## 计算每个特征对数据集的条件熵
             newEntropy = 0.0
             for value in values:
-                subDataSet = self._splitDataSet(dataSet, i, value)
-                prob = len(subDataSet)/float(len(dataSet))
-                newEntropy += prob*self._calcShannonEnt(subDataSet)
+                subDataArr, subClassList = self._splitDataSet(dataArr, classList, i, value)
+                prob = len(subClassList)/float(len(classList))
+                newEntropy += prob*self._calcShannonEnt(subClassList)
             '''
             信息增益 = 熵-条件熵
             '''
@@ -66,33 +66,34 @@ def _chooseBestFeatureToSplit(self, dataSet):
 
     def _majorityCnt(self, classList):
         classCount = {}
-        for vote in classList:
-            if vote not in classCount:
-                classCount[vote] = 1
+        for currentLabel in classList:
+            if currentLabel not in classCount:
+                classCount[currentLabel] = 1
             else:
-                classCount[vote] += 1
-            # if vote not in classCount:
-            #     classCount[vote] = 0
-            # classCount[vote] += 1
+                classCount[currentLabel] += 1
+            # if currentLabel not in classCount:
+            #     classCount[currentLabel] = 0
+            # classCount[currentLabel] += 1
         sortedClassCount = sorted(classCount.items(), key=lambda xx:xx[1], reverse=True)
         return sortedClassCount[0][0]
 
-    def fit(self, dataSet, featLabels):
-        classList = [data[-1] for data in dataSet]
+    def fit(self, dataArr, classList, featLabels):
         if classList.count(classList[0]) == len(classList):
             return classList[0] # 所有的类标签都相同，则返回类标签
-        if len(dataSet[0]) == 1: # 所有的类标签不完全相同，但用完所有特征，则返回次数最多的类标签
+        if len(dataArr[0]) == 0: # 所有的类标签不完全相同，但用完所有特征，则返回次数最多的类标签
             return self._majorityCnt(classList)
-        bestFeat = self._chooseBestFeatureToSplit(dataSet)
+        bestFeat = self._chooseBestFeatureToSplit(dataArr, classList)
         bestFeatLabel = featLabels[bestFeat]
         tree = {bestFeatLabel:{}}
         featLabels_copy = featLabels[:] # 这样不会改变输入的featLabels
         featLabels_copy.remove(bestFeatLabel)
-        featList = [data[bestFeat] for data in dataSet]
+        featList = [data[bestFeat] for data in dataArr]
         values = set(featList)
         for value in values:
-            subfeatLabels_copy = featLabels_copy[:] # 列表复制，非列表引用
-            tree[bestFeatLabel][value] = self.fit(self._splitDataSet(dataSet, bestFeat, value), subfeatLabels_copy)
+            subFeatLabels_copy = featLabels_copy[:] # 列表复制，非列表引用
+            subDataArr = self._splitDataSet(dataArr, classList, bestFeat, value)[0]
+            subClassList = self._splitDataSet(dataArr, classList, bestFeat, value)[1]
+            tree[bestFeatLabel][value] = self.fit(subDataArr, subClassList, subFeatLabels_copy)
         return tree
 
     def predict(self, tree, featLabels, testVec):
@@ -113,14 +114,19 @@ def loadDataSet():
                [1, 0, 'no'],
                [0, 1, 'no'],
                [0, 1, 'no']]
+    featData = []
+    labelData = []
+    for data in dataSet:
+        featData.append(data[:-1])
+        labelData.append(data[-1])
     featLabels = ['no surfacing', 'flippers'] # 特征标签
-    return dataSet, featLabels
+    return featData, labelData, featLabels
 
 if __name__ == '__main__':
-    myDataSet, myFeatLabels = loadDataSet()
-    print myDataSet, myFeatLabels
+    myFeatData, myLabelData, myFeatLabels = loadDataSet()
+    print myFeatData, myLabelData, myFeatLabels
     dt = DesicionTree()
-    myTree = dt.fit(myDataSet, myFeatLabels)
+    myTree = dt.fit(myFeatData, myLabelData, myFeatLabels)
     print myTree
     results = dt.predict(myTree, myFeatLabels, [1, 1])
     print results