机器学习---预测电影类型

#coding=UTF8 
from numpy import *
import operator

def createDataSet():
    group = array([[3,104],[2,100],[1,81],[101,10],[99,5]]) #训练集
    labels = ['affectional film','affectional film','affectional film','action movie',"action movie"]
    return group,labels

def classify0(inX,dataSet,labels,k):
    dataSetSize = dataSet.shape[0]
    diffMat = tile(inX, (dataSetSize,1)) - dataSet
    sqDiffMat = diffMat ** 2; #diffMat ^ 2
    sqDistances = sqDiffMat.sum(axis=1) #将矩阵的每一行相加比如[[2,1,3],[1,1,1]]结果为[6,3]
    distances = sqDistances ** 0.5 #sqDistances ^ (1/2)
    sortedDistIndicies = distances.argsort()
    classCount = {}
    for i in range(k):
        voteIlabel = labels[sortedDistIndicies[i]]
        classCount[voteIlabel] = classCount.get(voteIlabel,0) + 1
    sortedClassCount = sorted(classCount.iteritems(),key=operator.itemgetter(1),reverse=True)   #获取对象的第1个域的值
    return sortedClassCount[0][0]

group,labels = createDataSet()
print("输入武打镜头的个数:")
fun1 = input() #武打镜头的个数
print("输入亲吻镜头的个数:")
fun2 = input() #亲吻镜头的个数
b = classify0([fun1,fun2], group, labels, 3) #[7,100],7代表武打镜头为7,100代表接吻镜头为100,根据训练集可以得出这是一部爱情片
if b == "affectional film" :
    print("这很可能是一部爱情片")
else:
    print("这很可能是一部动作片")