DecisionTree.py 文件源码

def build_decision_tree(filename):
    """
        ??????????????
    """
    f=open(sys.argv[1],'r')
    reader=csv.reader(f)
    x=[]
    y=[]
    for line in reader:
        if line[1] in ['1','2','3']:#??????,??????
            x.append(line[2:4]+line[5:])
            y.append(line[1])
    x_train,x_test,y_train,y_test=cross_validation.train_test_split(x,y, test_size=0.2, random_state=42)
    clf=tree.DecisionTreeClassifier(max_depth=5)
    clf=clf.fit(x_train,y_train)
    score=clf.score(x_test,y_test)
    print score
    return clf,score