# -*- coding: utf-8 -*-


from numpy import *
import math
import copy
import cPickle as pickle


class C45DTree(object):
 def __init__(self): # 構(gòu)造方法
  self.tree = {} # 生成樹(shù)
  self.dataSet = [] # 數(shù)據(jù)集
  self.labels = [] # 標(biāo)簽集


 # 數(shù)據(jù)導(dǎo)入函數(shù)
 def loadDataSet(self, path, labels):
  recordList = []
  fp = open(path, "rb") # 讀取文件內(nèi)容
  content = fp.read()
  fp.close()
  rowList = content.splitlines() # 按行轉(zhuǎn)換為一維表
  recordList = [row.split("\t") for row in rowList if row.strip()] # strip()函數(shù)刪除空格、Tab等
  self.dataSet = recordList
  self.labels = labels


 # 執(zhí)行決策樹(shù)函數(shù)
 def train(self):
  labels = copy.deepcopy(self.labels)
  self.tree = self.buildTree(self.dataSet, labels)


 # 構(gòu)件決策樹(shù)：穿件決策樹(shù)主程序
 def buildTree(self, dataSet, lables):
  cateList = [data[-1] for data in dataSet] # 抽取源數(shù)據(jù)集中的決策標(biāo)簽列
  # 程序終止條件1：如果classList只有一種決策標(biāo)簽，停止劃分，返回這個(gè)決策標(biāo)簽
  if cateList.count(cateList[0]) == len(cateList):
   return cateList[0]
  # 程序終止條件2：如果數(shù)據(jù)集的第一個(gè)決策標(biāo)簽只有一個(gè)，返回這個(gè)標(biāo)簽
  if len(dataSet[0]) == 1:
   return self.maxCate(cateList)
  # 核心部分
  bestFeat, featValueList= self.getBestFeat(dataSet) # 返回?cái)?shù)據(jù)集的最優(yōu)特征軸
  bestFeatLabel = lables[bestFeat]
  tree = {bestFeatLabel: {}}
  del (lables[bestFeat])
  for value in featValueList: # 決策樹(shù)遞歸生長(zhǎng)
   subLables = lables[:] # 將刪除后的特征類別集建立子類別集
   # 按最優(yōu)特征列和值分隔數(shù)據(jù)集
   splitDataset = self.splitDataSet(dataSet, bestFeat, value)
   subTree = self.buildTree(splitDataset, subLables) # 構(gòu)建子樹(shù)
   tree[bestFeatLabel][value] = subTree
  return tree


 # 計(jì)算出現(xiàn)次數(shù)最多的類別標(biāo)簽
 def maxCate(self, cateList):
  items = dict([(cateList.count(i), i) for i in cateList])
  return items[max(items.keys())]


 # 計(jì)算最優(yōu)特征
 def getBestFeat(self, dataSet):
  Num_Feats = len(dataSet[0][:-1])
  totality = len(dataSet)
  BaseEntropy = self.computeEntropy(dataSet)
  ConditionEntropy = []  # 初始化條件熵
  slpitInfo = [] # for C4.5,caculate gain ratio
  allFeatVList = []
  for f in xrange(Num_Feats):
   featList = [example[f] for example in dataSet]
   [splitI, featureValueList] = self.computeSplitInfo(featList)
   allFeatVList.append(featureValueList)
   slpitInfo.append(splitI)
   resultGain = 0.0
   for value in featureValueList:
    subSet = self.splitDataSet(dataSet, f, value)
    appearNum = float(len(subSet))
    subEntropy = self.computeEntropy(subSet)
    resultGain += (appearNum/totality)*subEntropy
   ConditionEntropy.append(resultGain) # 總條件熵
  infoGainArray = BaseEntropy*ones(Num_Feats)-array(ConditionEntropy)
  infoGainRatio = infoGainArray/array(slpitInfo) # C4.5信息增益的計(jì)算
  bestFeatureIndex = argsort(-infoGainRatio)[0]
  return bestFeatureIndex, allFeatVList[bestFeatureIndex]

 # 計(jì)算劃分信息
 def computeSplitInfo(self, featureVList):
  numEntries = len(featureVList)
  featureVauleSetList = list(set(featureVList))
  valueCounts = [featureVList.count(featVec) for featVec in featureVauleSetList]
  pList = [float(item)/numEntries for item in valueCounts]
  lList = [item*math.log(item, 2) for item in pList]
  splitInfo = -sum(lList)
  return splitInfo, featureVauleSetList


 # 計(jì)算信息熵
 # @staticmethod
 def computeEntropy(self, dataSet):
  dataLen = float(len(dataSet))
  cateList = [data[-1] for data in dataSet] # 從數(shù)據(jù)集中得到類別標(biāo)簽
  # 得到類別為key、 出現(xiàn)次數(shù)value的字典
  items = dict([(i, cateList.count(i)) for i in cateList])
  infoEntropy = 0.0
  for key in items: # 香農(nóng)熵： = -p*log2(p) --infoEntropy = -prob * log(prob, 2)
   prob = float(items[key]) / dataLen
   infoEntropy -= prob * math.log(prob, 2)
  return infoEntropy


 # 劃分?jǐn)?shù)據(jù)集： 分割數(shù)據(jù)集； 刪除特征軸所在的數(shù)據(jù)列，返回剩余的數(shù)據(jù)集
 # dataSet : 數(shù)據(jù)集； axis: 特征軸； value: 特征軸的取值
 def splitDataSet(self, dataSet, axis, value):
  rtnList = []
  for featVec in dataSet:
   if featVec[axis] == value:
    rFeatVec = featVec[:axis] # list操作：提取0~（axis-1）的元素
    rFeatVec.extend(featVec[axis + 1:]) # 將特征軸之后的元素加回
    rtnList.append(rFeatVec)
  return rtnList

 # 存取樹(shù)到文件
 def storetree(self, inputTree, filename):
  fw = open(filename,'w')
  pickle.dump(inputTree, fw)
  fw.close()

 # 從文件抓取樹(shù)
 def grabTree(self, filename):
  fr = open(filename)
  return pickle.load(fr)

# -*- coding: utf-8 -*-

from numpy import *
from C45DTree import *

dtree = C45DTree()
dtree.loadDataSet("dataset.dat",["age", "revenue", "student", "credit"])
dtree.train()

dtree.storetree(dtree.tree, "data.tree")
mytree = dtree.grabTree("data.tree")
print mytree