PTA K近邻 python

我们学习了K近邻算法。
题目的任务是找出K个最近的邻居。
这里处理的样本只有一个特征，特征的数据类型都是数值型的，距离采用欧式距离。

输入格式:

第一行，算法的参数K
第二行，测试样本的特征
第三行，训练样本集的特征，不同样本之间以空格间隔

输出格式:

按距离从小到大输出K个最近样本的编号，样本编号按照输出样本的顺序编号，从0开始。每个样本编号后跟一个空格

输入样例

在这里给出输入样例。例如：

3
1
1 2 3 4 5 6

输出样例:

在这里给出相应的输出。例如：

0 1 2 

def k_nearest_neighbors(K, test_sample, training_samples):
    # Calculate Euclidean distance
    distances = [(i, (sample - test_sample) ** 2) for i, sample in enumerate(training_samples)]
    
    # Sort by distance
    sorted_distances = sorted(distances, key=lambda x: x[1])
    
    # Get K nearest neighbors
    neighbors = [x[0] for x in sorted_distances[:K]]
    
    return neighbors

# Test with the provided sample
K = int(inp