python之使用 matplotlib 连接散点图中 k 个最近邻居的有效方法

我正在尝试构建一个最近邻图，即散点图，其中每个数据点都连接到它的 k 个最近邻。我当前的解决方案有效，但显然效率不高。这是我到目前为止得到的:

import numpy as np 
from scipy.spatial.distance import pdist, squareform 
from matplotlib import pyplot as plt 
 
X = np.random.random(500).reshape((250, 2)) 
k = 4 
 
# matrix of pairwise Euclidean distances 
distmat = squareform(pdist(X, 'euclidean')) 
 
# select the kNN for each datapoint 
neighbors = np.sort(np.argsort(distmat, axis=1)[:, 0:k]) 
 
plt.figure(figsize = (8, 8)) 
plt.scatter(X[:,0], X[:,1], c = 'black') 
for i in np.arange(250): 
    for j in np.arange(k): 
        x1 = np.array([X[i,:][0], X[neighbors[i, j], :][0]]) 
        x2 = np.array([X[i,:][1], X[neighbors[i, j], :][1]]) 
        plt.plot(x1, x2, color = 'black') 
plt.show() 

有没有更有效的方法来构建这个图？

请您参考如下方法:

使用 LineCollection 一次绘制所有边，而不是单独绘制它们:

import numpy as np 
from scipy.spatial.distance import pdist, squareform 
from matplotlib import pyplot as plt 
from matplotlib.collections import LineCollection 
 
N = 250 
X = np.random.rand(250,2) 
k = 4 
 
# matrix of pairwise Euclidean distances 
distmat = squareform(pdist(X, 'euclidean')) 
 
# select the kNN for each datapoint 
neighbors = np.sort(np.argsort(distmat, axis=1)[:, 0:k]) 
 
# get edge coordinates 
coordinates = np.zeros((N, k, 2, 2)) 
for i in np.arange(250): 
    for j in np.arange(k): 
        coordinates[i, j, :, 0] = np.array([X[i,:][0], X[neighbors[i, j], :][0]]) 
        coordinates[i, j, :, 1] = np.array([X[i,:][1], X[neighbors[i, j], :][1]]) 
 
# create line artists 
lines = LineCollection(coordinates.reshape((N*k, 2, 2)), color='black') 
 
fig, ax = plt.subplots(1,1,figsize = (8, 8)) 
ax.scatter(X[:,0], X[:,1], c = 'black') 
ax.add_artist(lines) 
plt.show() 

在我的机器上，你的代码运行大约需要 1 秒；我的版本需要 65 毫秒。