作者：小小明
10年編碼經驗，熟悉Java、Python和Scala，非常擅長解決各類複雜數據處理的邏輯，各類結構化與非結構化數據互轉，字符串解析匹配等等。
至今已經幫助至少百名數據從業者解決工作中的實際問題，如果你在數據處理上遇到什麼困難，歡迎評論區與我交流。
求每個學生分數最接近的10個學生
需求背景
某MOOC課程網站的老師需要統計每個學生成績最相近的10個學生，距離計算公式是每門課程的成績之差的絕對值求和。
例如，學生1的成績爲(83,84,86,99,87)，學生2的成績爲(83,84,86,99,87)
兩個學生的距離爲|83-83|+|84-84|+|86-86|+|99-99|+|87-87|=0則判定爲這兩個學生成績距離完全相同，這個距離越大說明兩個學生的成績差異越大。
爲了保護學生隱私，下面僅使用模擬數據展示。
首先讀取數據：
import pandas as pd
import numpy as np
import heapq
df = pd.read_csv("學生成績.csv", encoding="gbk")
df.head()
結果：
這位老師管理的小班，僅100個學生。
採用雙重for循環遍歷每個學生與其他99個學生進行匹配計算，也僅1萬次循環的計算量，下面我將採用雙重for循環進行笛卡爾積計算，遍歷出每個學生分數最接近的10個學生。
笛卡爾積+最小堆解決需求
data = df.values
length = data.shape[0]
result = []
for i in range(length):
# 遍歷取出每一行的數據
src = data[i]
# 分別取出姓名和成績列表
src_name, src_scores = src[0], src[1:]
# 用一個最小堆保存最接近的10學生
min_similar_10 = []
for j in range(length):
if i == j:
# 跳過對自己的比較
continue
find = data[j]
# 被比較的學生姓名和成績
find_name, find_scores = find[0], find[1:]
# 計算兩個學生成績的距離
sim_value = np.abs(src_scores-find_scores).sum()
# 將當前學習和最小距離保存到最小堆中
heapq.heappush(min_similar_10, (find_name, sim_value))
if len(min_similar_10) > 10:
# 只保留10個距離最小的學生
min_similar_10 = heapq.nsmallest(
10, min_similar_10, key=lambda x: x[1])
name_similars, distances = list(zip(*min_similar_10))
result.append((src_name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
result
最終結果：
僅耗時200毫秒以內，說明笛卡爾積在100數據量時，可以順利解決這個問題的。
下面簡單介紹一下最小堆：
堆是數據結構中最常見的一種數據結構，是一個完全二叉樹。最小堆中每一個節點的值都小於等於其子樹中每個節點的值。
具體的原理，學過數據結構的讀者都懂，沒有學過的也不用深究，只需要知道它能夠實現快速找到N個最小值。
基於最小堆去保留最小的N個數據，遠比排序後取前N個快，幾乎就是logN算法複雜度和nlogN算法複雜度的差異。
官方文檔：https://docs.python.org/zh-cn/3/library/heapq.html?highlight=heapq#module-heapq
實現源碼：https://github.com/python/cpython/blob/3.9/Lib/heapq.py
如果你確實對堆的實現原理很感興趣，可以參考我的數據結構學習筆記：
https://datastructure.xiaoxiaoming.xyz/#/16.%E5%A0%86
使用numpy向量化操作解決需求
好景不長，該MOOC網站的管理員覺得我的代碼處理的效果比較好，希望把歷史班級所有學生的分數最接近的10個學生都找出來，大概有1W以上的學生。
用我上面的代碼跑了一下，結果跑了10分鐘也沒能計算完，問我能不能優化一下代碼。其實10分鐘跑不完太正常了，之前100個學生其實只需要循環大概1萬次，而這次1萬個學生卻需要循環大概1億次，耗時差距接近1萬倍，當然如果他再多等等，20分鐘也能跑完。
不過我個人也無法容忍要等那麼久，經過我一番代碼優化，程序在2分鐘內對1萬條數據跑出了結果。
考慮到部分很多讀者都沒有看懂上面的操作，這次我將用一個簡單的方法分步拆解，而且代碼更短，性能更高。
直接對本文開頭的100條數據進行測試，首先獲取源數據的姓名和分數數組：
data = df.values
length = data.shape[0]
names = data[:, 0]
scores = data[:, 1:]
print(names[:5])
print(scores[:5])
結果：
['學生1' '學生2' '學生3' '學生4' '學生5']
[[71 82 81 84 73 95 62 96 96 87 61 78 79 89 98 80]
[57 91 99 74 93 93 58 74 89 84 63 62 78 57 94 74]
[72 59 76 60 99 71 73 74 72 74 85 79 100 88 80 91]
[71 62 81 80 57 84 92 95 63 59 84 64 90 58 60 75]
[83 62 95 76 57 75 97 57 65 63 84 73 74 59 62 93]]
測試對第一個學生求分數最接近的10個學生：
for i in range(length):
name, score = names[i], scores[i]
print(name, score)
break
結果：
學生1 [77 89 60 99 83 80 95 93 94 82 92 93 96 98 70 84]
利用numpy向量化操作一次性求出該學生與其他所有學生的距離：
score_diff = np.abs(scores-score).sum(axis=1)
score_diff
結果：
array([0, 322, 248, 203, 215, 366, 198, 224, 299, 229, 274, 240, 253, 238,
183, 239, 223, 344, 320, 221, 291, 264, 273, 309, 307, 243, 258,
231, 225, 185, 234, 324, 315, 247, 247, 232, 179, 210, 238, 261,
175, 252, 259, 254, 330, 270, 193, 248, 252, 214, 175, 219, 262,
256, 265, 249, 263, 180, 270, 294, 260, 195, 272, 240, 213, 181,
252, 292, 261, 233, 238, 205, 237, 241, 244, 234, 255, 151, 228,
241, 173, 174, 221, 248, 191, 249, 243, 224, 252, 266, 229, 231,
228, 346, 232, 273, 269, 336, 294, 277], dtype=object)
取出距離最短的11個學生的索引，然後刪除自身的索引：
min_similar_index = np.argpartition(score_diff, 11)[:11].tolist()
min_similar_index.remove(i)
min_similar_index
結果：
[77, 80, 81, 40, 50, 36, 57, 65, 14, 29]
根據這10個學生的索引讀取所需要的數據：
name_similars = names[min_similar_index]
distances = score_diff[min_similar_index]
print(name_similars)
print(distances)
結果：
['學生78' '學生81' '學生82' '學生41' '學生51' '學生37' '學生58' '學生66' '學生15' '學生30']
[151 173 174 175 175 179 180 181 183 185]
這樣我們就已經計算出該學生成績最解決的10個學生的姓名和距離。
下面我整理一下完整處理代碼：
data = df.values
length = data.shape[0]
names = data[:, 0]
scores = data[:, 1:]
result = []
for i in range(length):
name, score = names[i], scores[i]
score_diff = np.abs(scores-score).sum(axis=1)
min_similar_index = np.argpartition(score_diff, 11)[:11].tolist()
min_similar_index.remove(i)
name_similars = names[min_similar_index]
distances = score_diff[min_similar_index]
result.append((name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
result
結果：
下面再針對一萬條數據跑一跑：
df = pd.read_csv("學生成績_10000.csv", encoding="gbk")
data = df.values
length = data.shape[0]
names = data[:, 0]
scores = data[:, 1:]
result = []
for i in range(length):
name, score = names[i], scores[i]
score_diff = np.abs(scores-score).sum(axis=1)
min_similar_index = np.argpartition(score_diff, 11)[:11].tolist()
min_similar_index.remove(i)
name_similars = names[min_similar_index]
distances = score_diff[min_similar_index]
result.append((name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
result
結果：
可以看到耗時爲1分48秒。
使用ball_tree解決需求
雖然上面的優化已經大幅度提升了程序性能，但畢竟仍然是O(n^2)算法複雜度的方法，萬一哪天網站要求對10萬個學生計算呢？時間又要多翻了接近幾十倍，耗時可能達到好幾個小時。我轉而一想直接用KNN內部的ball_tree來解決這個問題吧。
（關於使用KNN查找最近點的問題可參考很早之前的一篇文章：https://blog.csdn.net/as604049322/article/details/112385553）
先使用1000條學生成績的數據進行測試，讀取這1000條學生成績數據：
df = pd.read_csv("學生成績_1000.csv", encoding="gbk")
df
結果：
然後我們取出需要被訓練的數據：
# 取出用於被KNN訓練的數據
data = df.iloc[:, 1:].values
# y本身用於標註每條數據屬於哪個類別，但我並不使用KNN的分類功能，所以統一全部標註爲類別0
y = np.zeros(data.shape[0], dtype='int8')
print(data[:5])
print(y[:5])
結果：
[[ 77 89 60 99 83 80 95 93 94 82 92 93 96 98 70 84]
[ 90 63 72 57 100 85 63 62 67 77 82 68 81 79 100 71]
[ 71 71 82 72 62 70 82 63 92 69 97 97 82 69 63 57]
[ 88 78 88 61 72 61 77 93 94 88 86 95 84 70 73 94]
[ 66 67 59 70 75 67 79 74 57 84 96 67 93 87 58 83]]
[0 0 0 0 0]
創建KNN訓練器，並進行訓練：
from sklearn.neighbors import KNeighborsClassifier
knn = KNeighborsClassifier(n_neighbors=1, algorithm='ball_tree', p=1)
knn.fit(data, y)
distance, similar_points = knn.kneighbors(
data, n_neighbors=11, return_distance=True)
distance = distance.astype("int", copy=False)
print(distance[:5])
print(similar_points[:5])
結果：
這個過程的耗時僅50毫秒，幾乎可以忽略不計。
n_neighbors是KNN用來分類的參數，我並不使用它，將其指定的越小，越能減少無用的計算量，但是必須比0大，所以我指定爲1。
而需求方要求的距離計算公式顯然就等價於曼哈頓距離，所以我將p指定爲1，就跟需求方要求的距離計算公式一致。
knn.kneighbors則用來計算最近的點，n_neighbors指定爲11是因爲結果會包含自身，我打算後面再去除。
由於sklearn最終計算出來的距離是float浮點數類型，而我們的需求只可能產生整數距離，所以我將其轉換爲整數。
上面其實就相當於已經計算出了結果，下面我再將結果整理成需要的格式即可：
names = df['姓名']
result = []
for i, name in names.iteritems():
name_similar_indexs = similar_points[i].tolist()
self_index = name_similar_indexs.index(i)
name_similar_indexs.pop(self_index)
name_similars = names[name_similar_indexs].tolist()
distances = distance[i].tolist()
distances.pop(self_index)
result.append((name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
result
結果：
對於1000條數據ball_tree的計算耗時爲毫秒級。
下面我們對一萬條學生成績數據進行計算，首先讀取數據：
df = pd.read_csv("學生成績_10000.csv", encoding="gbk")
df
結果：
完整計算代碼：
# 取出用於被KNN訓練的數據
from sklearn.neighbors import KNeighborsClassifier
data = df.iloc[:, 1:].values
# y本身用於標註每條數據屬於哪個類別，但我並不使用KNN的分類功能，所以統一全部標註爲類別0
y = np.zeros(data.shape[0], dtype='int8')
knn = KNeighborsClassifier(n_neighbors=1, algorithm='ball_tree', p=1)
knn.fit(data, y)
distance, similar_points = knn.kneighbors(
data, n_neighbors=11, return_distance=True)
distance = distance.astype("int", copy=False)
names = df['姓名']
result = []
for i, name in names.iteritems():
name_similar_indexs = similar_points[i].tolist()
self_index = name_similar_indexs.index(i)
name_similar_indexs.pop(self_index)
name_similars = names[name_similar_indexs].tolist()
distances = distance[i].tolist()
distances.pop(self_index)
result.append((name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
result
結果：
可以看到1萬條數據ball_tree僅僅5秒就計算完了。
對4萬條學生數據測試
爲了測試方便，不再用excel來生成數據，而是直接使用python。
python直接生成測試數據的方法，以生成10條數據爲例：
size = 100000
data = np.c_[np.arange(1, size+1).reshape((-1, 1)),
np.random.randint(56, 101, size=(size, 16))]
df = pd.DataFrame(data, columns=["姓名", "Python", "C/C++", "Java", "Scala", "數據結構", "離散數學",
"計算機體系結構", "編譯原理", "計算機網絡", "數據庫原理", "計算機圖形學",
"自然語言處理", "嵌入式系統及應用", "網絡信息與安全", "計算機視覺", "人工智能"])
df["姓名"] = "學生"+df["姓名"].astype(str)
df
結果：
然後使用上面相同的代碼分別測試1w條、2w條、3w條、…、10w條：
1萬條耗時5.3秒。
2萬條耗時15.8秒。
3萬條耗時32.3秒。
4萬條耗時60.75秒。
預估10萬條數據的耗時
這不行，僅4萬條耗時就達到一分鐘，這也開始讓我有點等的捉急了，不能繼續測試下去了。
這時間增長趨勢好像也不是線性增長而是指數增長，下面先就記錄一下多少條記錄時耗時多久吧：
import pandas as pd
import numpy as np
import time
from sklearn.neighbors import KNeighborsClassifier
times = {
}
for size in np.r_[np.arange(1000, 10001, 1000), np.arange(20000, 40001, 10000)]:
data = np.c_[np.arange(1, size+1).reshape((-1, 1)),
np.random.randint(56, 101, size=(size, 16))]
df = pd.DataFrame(data, columns=["姓名", "Python", "C/C++", "Java", "Scala", "數據結構", "離散數學",
"計算機體系結構", "編譯原理", "計算機網絡", "數據庫原理", "計算機圖形學",
"自然語言處理", "嵌入式系統及應用", "網絡信息與安全", "計算機視覺", "人工智能"])
df["姓名"] = "學生"+df["姓名"].astype(str)
start_time = time.perf_counter()
# 取出用於被KNN訓練的數據
data = df.iloc[:, 1:].values
# y本身用於標註每條數據屬於哪個類別，但我並不使用KNN的分類功能，所以統一全部標註爲類別0
y = np.zeros(data.shape[0], dtype='int8')
knn = KNeighborsClassifier(n_neighbors=1, algorithm='ball_tree', p=1)
knn.fit(data, y)
distance, similar_points = knn.kneighbors(
data, n_neighbors=11, return_distance=True)
distance = distance.astype("int", copy=False)
names = df['姓名']
result = []
for i, name in names.iteritems():
name_similar_indexs = similar_points[i].tolist()
self_index = name_similar_indexs.index(i)
name_similar_indexs.pop(self_index)
name_similars = names[name_similar_indexs].tolist()
distances = distance[i].tolist()
distances.pop(self_index)
result.append((name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
take_time = time.perf_counter()-start_time
print(f"{size}條數據耗時{take_time:.2f}秒")
times[size] = take_time
time_df = pd.DataFrame.from_dict(times, orient='index', columns=["time"])
time_df.plot()
結果：
1000條數據耗時0.28秒
2000條數據耗時0.58秒
3000條數據耗時0.94秒
4000條數據耗時1.39秒
5000條數據耗時1.84秒
6000條數據耗時2.38秒
7000條數據耗時3.03秒
8000條數據耗時3.66秒
9000條數據耗時4.34秒
10000條數據耗時5.04秒
20000條數據耗時15.43秒
30000條數據耗時31.44秒
40000條數據耗時64.27秒
從這走勢來看有點像二次函數或冪次函數，我們假設這是一個二次函數然後使用numpy擬合這條曲線，並預估10萬數據的耗時：
from numpy import polyfit, poly1d
import matplotlib.pyplot as plt
%matplotlib inline
x = np.arange(1000, 100001, 1000)
y = poly1d(polyfit(time_df.index, time_df.time, 2))
plt.figure(figsize=(10, 6))
plt.plot(time_df.index, time_df.time, 'rx')
plt.plot(x, y(x), 'b:')
plt.show()
print(f"10萬條數據預計耗時{y(100000):.2f}秒")
結果：
預計耗時6分鐘，但如果這個時候還用笛卡爾積去計算，估計耗時幾個小時。
不過實際測試了一下，好像10分鐘也沒有出結果，看來樣本還是太少，曲線擬合的效果還是不太準。
測試到6W試一下：
import pandas as pd
import numpy as np
import time
from sklearn.neighbors import KNeighborsClassifier
times = {
}
for size in np.r_[np.arange(1000, 10001, 1000), np.arange(20000, 60001, 10000)]:
data = np.c_[np.arange(1, size+1).reshape((-1, 1)),
np.random.randint(56, 101, size=(size, 16))]
df = pd.DataFrame(data, columns=["姓名", "Python", "C/C++", "Java", "Scala", "數據結構", "離散數學",
"計算機體系結構", "編譯原理", "計算機網絡", "數據庫原理", "計算機圖形學",
"自然語言處理", "嵌入式系統及應用", "網絡信息與安全", "計算機視覺", "人工智能"])
df["姓名"] = "學生"+df["姓名"].astype(str)
start_time = time.perf_counter()
# 取出用於被KNN訓練的數據
data = df.iloc[:, 1:].values
# y本身用於標註每條數據屬於哪個類別，但我並不使用KNN的分類功能，所以統一全部標註爲類別0
y = np.zeros(data.shape[0], dtype='int8')
knn = KNeighborsClassifier(n_neighbors=1, algorithm='ball_tree', p=1)
knn.fit(data, y)
distance, similar_points = knn.kneighbors(
data, n_neighbors=11, return_distance=True)
distance = distance.astype("int", copy=False)
names = df['姓名']
result = []
for i, name in names.iteritems():
name_similar_indexs = similar_points[i].tolist()
self_index = name_similar_indexs.index(i)
name_similar_indexs.pop(self_index)
name_similars = names[name_similar_indexs].tolist()
distances = distance[i].tolist()
distances.pop(self_index)
result.append((name, name_similars, distances))
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
take_time = time.perf_counter()-start_time
print(f"{size}條數據耗時{take_time:.2f}秒")
times[size] = take_time
time_df = pd.DataFrame.from_dict(times, orient='index', columns=["time"])
結果：
1000條數據耗時0.29秒
2000條數據耗時0.59秒
3000條數據耗時0.98秒
4000條數據耗時1.40秒
5000條數據耗時1.97秒
6000條數據耗時2.44秒
7000條數據耗時3.00秒
8000條數據耗時3.74秒
9000條數據耗時4.41秒
10000條數據耗時5.23秒
20000條數據耗時15.92秒
30000條數據耗時31.75秒
40000條數據耗時64.25秒
50000條數據耗時132.26秒
60000條數據耗時220.87秒
再擬合一下：
from numpy import polyfit, poly1d
import matplotlib.pyplot as plt
%matplotlib inline
x = np.arange(1000, 100001, 1000)
y = poly1d(polyfit(time_df.index, time_df.time, 2))
plt.figure(figsize=(10, 6))
plt.plot(time_df.index, time_df.time, 'rx')
plt.plot(x, y(x), 'b:')
plt.show()
print(f"10萬條數據預計耗時{y(100000):.2f}秒")
結果：
從擬合效果來看，可能這個虛線實際並不是二次函數，而是3次以上的函數，下面使用3次函數進行擬合：
from numpy import polyfit, poly1d
import matplotlib.pyplot as plt
%matplotlib inline
x = np.arange(1000, 100001, 1000)
y = poly1d(polyfit(time_df.index, time_df.time, 3))
plt.figure(figsize=(10, 6))
plt.plot(time_df.index, time_df.time, 'rx')
plt.plot(x, y(x), 'b:')
plt.show()
print(f"10萬條數據預計耗時{y(100000):.2f}秒")
結果：
預估耗時爲17分鐘。
實際耗時呢？
import pandas as pd
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
size = 100000
data = np.c_[np.arange(1, size+1).reshape((-1, 1)),
np.random.randint(56, 101, size=(size, 16))]
df = pd.DataFrame(data, columns=["姓名", "Python", "C/C++", "Java", "Scala", "數據結構", "離散數學",
"計算機體系結構", "編譯原理", "計算機網絡", "數據庫原理", "計算機圖形學",
"自然語言處理", "嵌入式系統及應用", "網絡信息與安全", "計算機視覺", "人工智能"])
df["姓名"] = "學生"+df["姓名"].astype(str)
# 取出用於被KNN訓練的數據
data = df.iloc[:, 1:].values
# y本身用於標註每條數據屬於哪個類別，但我並不使用KNN的分類功能，所以統一全部標註爲類別0
y = np.zeros(data.shape[0], dtype='int8')
knn = KNeighborsClassifier(n_neighbors=1, algorithm='ball_tree', p=1)
knn.fit(data, y)
distance, similar_points = knn.kneighbors(
data, n_neighbors=11, return_distance=True)
distance = distance.astype("int", copy=False)
names = df['姓名']
結果：
訓練耗時13分鐘。
from tqdm.notebook import tqdm
result = []
pbar = tqdm(total=size)
for i, name in names.items():
pbar.update(1)
name_similar_indexs = similar_points[i].tolist()
self_index = name_similar_indexs.index(i)
name_similar_indexs.pop(self_index)
name_similars = names[name_similar_indexs].tolist()
distances = distance[i].tolist()
distances.pop(self_index)
result.append((name, name_similars, distances))
pbar.close()
result = pd.DataFrame(result, columns=["姓名", "分數最接近的10個學生", "距離"])
result
整理結果耗時半分鐘，實際耗時是14分鐘，跟預測的17分鐘差不多。
雖然14分鐘也比較慢，但相對前面的笛卡爾積的算法需要耗時好幾小時而言已經大幅度提升程序計算性能，節約了計算時間。
總結
今天我向你演示瞭如何使用最小堆和numpy來快速計算每個學生分數最接近的10個學生，可以看到在數據量小於10000時，這種時間複雜度爲O(n^2)的算法還可以接受，但一旦達到2萬以上，基本上就慢的難以忍受了。所以我使用ball_tree來計算這個距離，1W數據量僅耗時5秒。
但是當數據量達到5萬以上時，ball_tree也有點慢了，但也相對笛卡爾積的算法還是快了很多。
最後，我使用numpy擬合時間曲線，預估10w數據量時ball_tree的耗時爲17分鐘，實際測試是14分鐘，基本預測正確。
歡迎下方留言或評論，分享你的看法。