import pandas as pd
import numpy as np
 
#讀取數(shù)據(jù)集，至少包含變量和target兩列
sample_set = pd.read_excel('/數(shù)據(jù)樣本.xlsx')
 
def calc_score_median(sample_set, var):
  '''
  計(jì)算相鄰評分的中位數(shù)，以便進(jìn)行決策樹二元切分
  param sample_set: 待切分樣本
  param var: 分割變量名稱
  '''
  var_list = list(np.unique(sample_set[var]))
  var_median_list = []
  for i in range(len(var_list) -1):
    var_median = (var_list[i] + var_list[i+1]) / 2
    var_median_list.append(var_median)
  return var_median_list

def choose_best_split(sample_set, var, min_sample):
  '''
  使用CART分類決策樹選擇最好的樣本切分點(diǎn)
  返回切分點(diǎn)
  param sample_set: 待切分樣本
  param var: 分割變量名稱
  param min_sample: 待切分樣本的最小樣本量(限制條件)
  '''
  # 根據(jù)樣本評分計(jì)算相鄰不同分?jǐn)?shù)的中間值
  score_median_list = calc_score_median(sample_set, var)
  median_len = len(score_median_list)
  sample_cnt = sample_set.shape[0]
  sample1_cnt = sum(sample_set['target'])
  sample0_cnt = sample_cnt- sample1_cnt
  Gini = 1 - np.square(sample1_cnt / sample_cnt) - np.square(sample0_cnt / sample_cnt)
  
  bestGini = 0.0; bestSplit_point = 0.0; bestSplit_position = 0.0
  for i in range(median_len):
    left = sample_set[sample_set[var] < score_median_list[i]]
    right = sample_set[sample_set[var] > score_median_list[i]]
    
    left_cnt = left.shape[0]; right_cnt = right.shape[0]
    left1_cnt = sum(left['target']); right1_cnt = sum(right['target'])
    left0_cnt = left_cnt - left1_cnt; right0_cnt = right_cnt - right1_cnt
    left_ratio = left_cnt / sample_cnt; right_ratio = right_cnt / sample_cnt
    
    if left_cnt < min_sample or right_cnt < min_sample:
      continue
    
    Gini_left = 1 - np.square(left1_cnt / left_cnt) - np.square(left0_cnt / left_cnt)
    Gini_right = 1 - np.square(right1_cnt / right_cnt) - np.square(right0_cnt / right_cnt)
    Gini_temp = Gini - (left_ratio * Gini_left + right_ratio * Gini_right)
    if Gini_temp > bestGini:
      bestGini = Gini_temp; bestSplit_point = score_median_list[i]
      if median_len > 1:
        bestSplit_position = i / (median_len - 1)
      else:
        bestSplit_position = i / median_len
    else:
      continue
        
  Gini = Gini - bestGini
  return bestSplit_point, bestSplit_position

def bining_data_split(sample_set, var, min_sample, split_list):
  '''
  劃分?jǐn)?shù)據(jù)找到最優(yōu)分割點(diǎn)list
  param sample_set: 待切分樣本
  param var: 分割變量名稱
  param min_sample: 待切分樣本的最小樣本量(限制條件)
  param split_list: 最優(yōu)分割點(diǎn)list
  '''
  split, position = choose_best_split(sample_set, var, min_sample)
  if split != 0.0:
    split_list.append(split)
  # 根據(jù)分割點(diǎn)劃分?jǐn)?shù)據(jù)集，繼續(xù)進(jìn)行劃分
  sample_set_left = sample_set[sample_set[var] < split]
  sample_set_right = sample_set[sample_set[var] > split]
  # 如果左子樹樣本量超過2倍最小樣本量，且分割點(diǎn)不是第一個分割點(diǎn)，則切分左子樹
  if len(sample_set_left) >= min_sample * 2 and position not in [0.0, 1.0]:
    bining_data_split(sample_set_left, var, min_sample, split_list)
  else:
    None
  # 如果右子樹樣本量超過2倍最小樣本量，且分割點(diǎn)不是最后一個分割點(diǎn)，則切分右子樹
  if len(sample_set_right) >= min_sample * 2 and position not in [0.0, 1.0]:
    bining_data_split(sample_set_right, var, min_sample, split_list)
  else:
    None

def get_bestsplit_list(sample_set, var):
  '''
  根據(jù)分箱得到最優(yōu)分割點(diǎn)list
  param sample_set: 待切分樣本
  param var: 分割變量名稱
  '''
  # 計(jì)算最小樣本閾值（終止條件）
  min_df = sample_set.shape[0] * 0.05
  split_list = []
  # 計(jì)算第一個和最后一個分割點(diǎn)
  bining_data_split(sample_set, var, min_df, split_list)
  return split_list