[027量化交易] python knn实例

import numpy as np
import matplotlib.pyplot as plt
import operator
# https://blog.csdn.net/qq_62945476/article/details/123737155

'''
    trainData - 训练集
    testData - 测试集
    labels - 分类
'''


def knn(trainData, testData, labels, k):
    # 计算训练样本的行数
    rowSize = trainData.shape[0]
    # 计算训练样本和测试样本的差值
    diff = np.tile(testData, (rowSize, 1)) - trainData
    # 计算差值的平方和
    sqrDiff = diff ** 2
    sqrDiffSum = sqrDiff.sum(axis=1)
    # 计算距离
    distances = sqrDiffSum ** 0.5
    # 对所得的距离从低到高进行排序
    print(distances)
    sortDistance = distances.argsort()
    print(sortDistance)

    count = {}

    for i in range(k):
        vote = labels[sortDistance[i]]
        print('----------------------')
        print(sortDistance[i])
        print(vote)
        count[vote] = count.get(vote, 0) + 1
    # 对类别出现的频数从高到低进行排序
    sortCount = sorted(count.items(), key=operator.itemgetter(1), reverse=True)

    # 返回出现频数最高的类别
    return sortCount[0][0]


trainData = np.array([[5, 1], [4, 0], [1, 3], [0, 4]])
labels = ['动作片', '动作片', '爱情片', '爱情片']
print(trainData)
testData = [3, 2]
X = knn(trainData, testData, labels, 3)
print('**************************')
print(X)

def data_show(in_data, train_data):
    # 显示训练数据
    x = []
    y = []
    for i in range(train_data.shape[0]):
        x.append(train_data[i][0])
        y.append(train_data[i][1])
    plt.plot(x, y, "*")
    plt.xlabel("Number of kisses")
    plt.ylabel("Number of fights")
    plt.plot(in_data[0], in_data[1], "r*")
    plt.show()


data_show(testData,trainData)