2,Root Node Attribute Selection for Decision Trees using Information Gain

from sklearn.datasets import load_iris

from sklearn.tree import DecisionTreeClassifier


# Load the built-in Iris dataset

data = load_iris()

X = data.data

y = data.target

feature_names = data.feature_names


# Create a decision tree using entropy (information gain)

model = DecisionTreeClassifier(criterion='entropy', max_depth=1)  # Only root node

model.fit(X, y)


# Get the feature used at the root node

root_feature_index = model.tree_.feature[0]

root_feature_name = feature_names[root_feature_index]


print("Root Node Attribute (Best Feature):", root_feature_name)


Comments