I am working with the titanic dataset and using decision trees for analyzing the age covariate. I'd like just to see whether kids are more likely to survive than adults. I implemented my own Gini coefficient and I had plot the coefficient by age: dataset here titanic ds
import pandas as pd import seaborn import matplotlib.pyplot as plt from sklearn import tree import graphviz import numpy as np def gini_by_age(df, t): df['age_group'] = df['age'].apply(lambda row : 0 if row <= t else 1) kids = df[df['age_group'] == 0] kids0 = kids[kids['survived'] == 0] kids1 = kids[kids['survived'] == 1] adults = df[df['age_group'] == 1] adults0 = adults[adults['survived'] == 0] adults1 = adults[adults['survived'] == 1] gk = 1 - (len(kids0)**2 + len(kids1)**2)/float(len(kids))**2 ga = 1 - (len(adults0)**2 + len(adults1)**2)/float(len(adults))**2 return gk + ga def plot_gini_by_age(df): ages = range(2,25) y = [gini_by_age(df, a) for a in ages] plt.plot(ages, y) plt.show() def use_tree(df): X = np.array(df['age']).reshape((len(df['age']),1)) y = df['survived'] clf = tree.DecisionTreeClassifier(max_depth=1).fit(X,y) dot_data = tree.export_graphviz(clf, out_file=None) graph = graphviz.Source(dot_data) graph.render("age") titanic_df = pd.read_csv("titanic_ds.csv") ages_cov = titanic_df[['age', 'survived']].dropna() plot_gini_by_age(ages_cov) use_tree(ages_cov) print gini_by_age(ages_cov, 5) print gini_by_age(ages_cov, 8.5) print gini_by_age(ages_cov, 15)
output: 0.925844132419 0.937732003001 0.963875889772 I see from the plot that gini coefficient has local minima at roughly 5, 8 and 15 years old and the best is at 5. But scikit gives me 8.5 years old as the best split. What is wrong here?