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