# -*- coding: utf-8 -*-
__author__ = 'zhihua_oba'

import operator
from numpy import *
from math import log

#文件讀取
def file2matrix(filename, attribute_num): #傳入?yún)?shù)：文件名，屬性個數(shù)
 fr = open(filename)
 arrayOLines = fr.readlines()
 numberOfLines = len(arrayOLines) #統(tǒng)計數(shù)據(jù)集行數(shù)（樣本個數(shù)）
 dataMat = zeros((numberOfLines, attribute_num))
 classLabelVector = [] #分類標簽
 index = 0
 for line in arrayOLines:
  line = line.strip() #strip() 刪除字符串中的'\n'
  listFromLine = line.split() #將一個字符串分裂成多個字符串組成的列表，不帶參數(shù)時以空格進行分割，當代參數(shù)時，以該參數(shù)進行分割
  dataMat[index, : ] = listFromLine[0:attribute_num] #讀取數(shù)據(jù)對象屬性值
  classLabelVector.append(listFromLine[-1]) #讀取分類信息
  index += 1
 dataSet = [] #數(shù)組轉化成列表
 index = 0
 for index in range(0, numberOfLines):
  temp = list(dataMat[index, :])
  temp.append(classLabelVector[index])
  dataSet.append(temp)
 return dataSet

#劃分數(shù)據(jù)集
def splitDataSet(dataSet, axis, value):
 retDataSet = []
 for featvec in dataSet: #每行
  if featvec[axis] == value: #每行中第axis個元素和value相等 #刪除對應的元素，并將此行，加入到rerDataSet
   reducedFeatVec = featvec[:axis]
   reducedFeatVec.extend(featvec[axis+1:])
   retDataSet.append(reducedFeatVec)
 return retDataSet

#計算香農(nóng)熵 #計算數(shù)據(jù)集的香農(nóng)熵 == 計算數(shù)據(jù)集類標簽的香農(nóng)熵
def calcShannonEnt(dataSet):
 numEntries = len(dataSet) #數(shù)據(jù)集樣本點個數(shù)
 labelCounts = {} #類標簽
 for featVec in dataSet: #統(tǒng)計數(shù)據(jù)集類標簽的個數(shù)，字典形式
  currentLabel = featVec[-1]
  if currentLabel not in labelCounts.keys():
   labelCounts[currentLabel] = 0
  labelCounts[currentLabel] += 1
 shannonEnt = 0.0
 for key in labelCounts:
  prob = float(labelCounts[key])/numEntries
  shannonEnt -= prob * log(prob, 2)
 return shannonEnt

#根據(jù)香農(nóng)熵，選擇最優(yōu)的劃分方式 #根據(jù)某一屬性劃分后，類標簽香農(nóng)熵越低，效果越好
def chooseBestFeatureToSplit(dataSet):
 baseEntropy = calcShannonEnt(dataSet) #計算數(shù)據(jù)集的香農(nóng)熵
 numFeatures = len(dataSet[0])-1
 bestInfoGain = 0.0 #最大信息增益
 bestFeature = 0 #最優(yōu)特征
 for i in range(0, numFeatures):
  featList = [example[i] for example in dataSet] #所有子列表（每行）的第i個元素，組成一個新的列表
  uniqueVals = set(featList)
  newEntorpy = 0.0
  for value in uniqueVals: #數(shù)據(jù)集根據(jù)第i個屬性進行劃分，計算劃分后數(shù)據(jù)集的香農(nóng)熵
   subDataSet = splitDataSet(dataSet, i, value)
   prob = len(subDataSet)/float(len(dataSet))
   newEntorpy += prob*calcShannonEnt(subDataSet)
  infoGain = baseEntropy-newEntorpy #劃分后的數(shù)據(jù)集，香農(nóng)熵越小越好，即信息增益越大越好
  if(infoGain > bestInfoGain):
   bestInfoGain = infoGain
   bestFeature = i
 return bestFeature

#如果數(shù)據(jù)集已經(jīng)處理了所有屬性，但葉子結點中類標簽依然不是唯一的，此時需要決定如何定義該葉子結點。這種情況下，采用多數(shù)表決方法，對該葉子結點進行分類
def majorityCnt(classList): #傳入?yún)?shù)：葉子結點中的類標簽
 classCount = {}
 for vote in classList:
  if vote not in classCount.keys():
   classCount[vote] = 0
   classCount[vote] += 1
 sortedClassCount = sorted(classCount.iteritems(), key=operator.itemgetter(1), reverse=True)
 return sortedClassCount[0][0]

#創(chuàng)建樹
def createTree(dataSet, labels): #傳入?yún)?shù)：數(shù)據(jù)集，屬性標簽（屬性標簽作用：在輸出結果時，決策樹的構建更加清晰）
 classList = [example[-1] for example in dataSet] #數(shù)據(jù)集樣本的類標簽
 if classList.count(classList[0]) == len(classList): #如果數(shù)據(jù)集樣本屬于同一類，說明該葉子結點劃分完畢
  return classList[0]
 if len(dataSet[0]) == 1: #如果數(shù)據(jù)集樣本只有一列（該列是類標簽），說明所有屬性都劃分完畢，則根據(jù)多數(shù)表決方法，對該葉子結點進行分類
  return majorityCnt(classList)
 bestFeat = chooseBestFeatureToSplit(dataSet) #根據(jù)香農(nóng)熵，選擇最優(yōu)的劃分方式
 bestFeatLabel = labels[bestFeat] #記錄該屬性標簽
 myTree = {bestFeatLabel:{}} #樹
 del(labels[bestFeat]) #在屬性標簽中刪除該屬性
 #根據(jù)最優(yōu)屬性構建樹
 featValues = [example[bestFeat] for example in dataSet]
 uniqueVals = set(featValues)
 for value in uniqueVals:
  subLabels = labels[:]
  subDataSet = splitDataSet(dataSet, bestFeat, value)
  myTree[bestFeatLabel][value] = createTree(subDataSet, subLabels)
 return myTree

#測試算法：使用決策樹，對待分類樣本進行分類
def classify(inputTree, featLabels, testVec): #傳入?yún)?shù)：決策樹，屬性標簽，待分類樣本
 firstStr = inputTree.keys()[0] #樹根代表的屬性
 secondDict = inputTree[firstStr]
 featIndex = featLabels.index(firstStr) #樹根代表的屬性，所在屬性標簽中的位置，即第幾個屬性
 for key in secondDict.keys():
  if testVec[featIndex] == key:
   if type(secondDict[key]).__name__ == 'dict':
    classLabel = classify(secondDict[key], featLabels, testVec)
   else:
    classLabel = secondDict[key]
 return classLabel

def main():
 dataSet = file2matrix('test_sample.txt', 4)
 labels = ['attr01', 'attr02', 'attr03', 'attr04']
 labelsForCreateTree = labels[:]
 Tree = createTree(dataSet, labelsForCreateTree )
 testvec = [2, 3, 2, 3]
 print classify(Tree, labels, testvec)
if __name__ == '__main__':
  main()