gini.py

import sklearn.tree
import sklearn.preprocessing
import matplotlib.pyplot as plt
import pandas as pd

df = pd.read_csv('data/movie.csv')
labelEncoder = sklearn.preprocessing.LabelEncoder()
fig, axes = plt.subplots(1, 2, layout='tight')
for i, name in enumerate(('popcorn', 'coke')):
    decisionTreeClassifier = sklearn.tree.DecisionTreeClassifier(
        criterion = 'gini',
        random_state = 42,
    )
    df_encoded = pd.DataFrame()
    df_encoded[name] = labelEncoder.fit_transform(df[name])
    decisionTreeClassifier.fit(X=df_encoded, y=df['watch'])
    sklearn.tree.plot_tree(
        decisionTreeClassifier,
        ax            = axes[i],
        feature_names = [name],
        filled        = True,
        class_names   = ['naah', 'watch'],
        rounded       = True,
        fontsize      = 14,
        node_ids      = True,
        # proportion    = True,
        # impurity      = False,
    )
plt.show()
fig.savefig(f:='./py-sklearn-plt/movie-tree-1d.png'); return f