使用matplotlib

Question

我正在尝试构建一个最近邻图，即一个散点图，其中每个数据点连接到其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()

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

Answer 1

使用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毫秒。