from sklearn.neighbors import KNeighborsClassifier

# X中存储交通工具的参数
# 总长度（米）、时速（km/h）、重量（吨）、座位数量
X = [[96, 85, 120, 400],        # 普通火车
     [144, 92, 200, 600],
     [240, 87, 350, 1000],
     [360, 90, 495, 1300],
     [384, 91, 530, 1405],
     [240, 360, 490, 800],      # 高铁
     [360, 380, 750, 1200],
     [290, 380, 480, 960],
     [120, 320, 160, 400],
     [384, 340, 520, 1280],
     [33.4, 918, 77, 180],      # 飞机
     [33.6, 1120, 170.5, 185],
     [39.5, 785, 230, 240],
     [33.84, 940, 150, 195],
     [44.5, 920, 275, 275],
     [75.3, 1050, 575, 490]]
# y中存储类别，0表示普通火车，1表示高铁，2表示飞机
y = [0]*5+[1]*5+[2]*6
# labels中存储对应的交通工具名称
labels = ('普通火车', '高铁', '飞机')

# 创建并训练模型
knn = KNeighborsClassifier(n_neighbors=3,
                           weights='distance')
knn.fit(X, y)

# 对未知样本进行分类
unKnown = [[300, 79, 320, 900],
           [36.7, 800, 190, 220]]
result = knn.predict(unKnown)
for para, index in zip(unKnown, result):
    print(para, labels[index], sep=':')