# -*- coding:utf-8 -*-
#! python2
import numpy as np
from math import log
data_feature_matrix = np.array([[1, 1],
                [1, 1],
                [1, 0],
                [0, 1],
                [0, 1]]) # 特征矩陣
category = ['yes', 'yes', 'no', 'no', 'no'] # 5個對象分別所屬的類別
def calc_shannon_ent(category_list):
  """
  :param category_list: 類別列表
  :return: 該類別列表的熵值
  """
  label_count = {} # 統(tǒng)計數(shù)據(jù)集中每個類別的個數(shù)
  num = len(category_list) # 數(shù)據(jù)集個數(shù)
  for i in range(num):
    try:
      label_count[category_list[i]] += 1
    except KeyError:
      label_count[category_list[i]] = 1
  shannon_ent = 0.
  for k in label_count:
    prob = float(label_count[k]) / num
    shannon_ent -= prob * log(prob, 2) # 計算信息熵
  return shannon_ent
def split_data(feature_matrix, category_list, feature_index, value):
  """
  篩選出指定特征值所對應的類別列表
  :param category_list: 類別列表
  :param feature_matrix: 特征矩陣
  :param feature_index: 指定特征索引
  :param value: 指定特征屬性的特征值
  :return: 符合指定特征屬性的特征值的類別列表
  """
  # feature_matrix = np.array(feature_matrix)
  ret_index = np.where(feature_matrix[:, feature_index] == value)[0] # 獲取符合指定特征值的索引
  ret_category_list = [category_list[i] for i in ret_index] # 根據(jù)索引取得指定的所屬類別，構建為列表
  return ret_category_list
def choose_best_feature(feature_matrix, category_list):
  """
  根據(jù)信息增益獲取最優(yōu)特征
  :param feature_matrix: 特征矩陣
  :param category_list: 類別列表
  :return: 最優(yōu)特征對應的索引
  """
  feature_num = len(feature_matrix[0]) # 特征個數(shù)
  data_num = len(category_list) # 數(shù)據(jù)集的個數(shù)
  base_shannon_ent = calc_shannon_ent(category_list=category_list) # 原始數(shù)據(jù)的信息熵
  best_info_gain = 0 # 最優(yōu)信息增益
  best_feature_index = -1 # 最優(yōu)特征對應的索引
  for f in range(feature_num):
    uni_value_list = set(feature_matrix[:, f]) # 該特征屬性所包含的特征值
    new_shannon_ent = 0.
    for value in uni_value_list:
      sub_cate_list = split_data(feature_matrix=feature_matrix, category_list=category_list, feature_index=f, value=value)
      prob = float(len(sub_cate_list)) / data_num
      new_shannon_ent += prob * calc_shannon_ent(sub_cate_list)
    info_gain = base_shannon_ent - new_shannon_ent # 信息增益
    print '初始信息熵為：', base_shannon_ent, '按照特征%i分類后的信息熵為：' % f, new_shannon_ent, '信息增益為：', info_gain
    if info_gain > best_info_gain:
      best_info_gain = info_gain
      best_feature_index = f
  return best_feature_index
if __name__ == '__main__':
  best_feature = choose_best_feature(data_feature_matrix, category)
  print '最好用于劃分數(shù)據(jù)集的特征為：', best_feature