实现决策树算法

from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
 #决策树算法
def DataSet():
    group = np.array([[3, 104], [2, 100], [1, 81], [101, 10], [99, 5], [98, 2],[90,98]])
    labels = ['爱情电影', '爱情电影', '爱情电影', '动作电影', '动作电影', '动作电影','爱情电影']
    return group, labels

def DecisionTreeModel(in_x, traindata, labels):
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(traindata, labels, test_size=0.2, random_state=42)
    
    # 初始化决策树分类器
    clf = DecisionTreeClassifier()
    
    # 拟合模型
    clf.fit(X_train, y_train)
    
    # 预测
    y_pred = clf.predict([in_x])
    
    return y_pred[0]

if __name__ == '__main__':
    group, labels = DataSet()
    testdata = [90,98]
    result = DecisionTreeModel(testdata, group, labels)
    print('输入对应的数据所属类别是：{}'.format(result))