如何仅用NumPy复现Scikit-learn中kneighbors_graph的功能并理解其底层原理？

Got it, let's break down how scikit-learn's kneighbors_graph works and replicate its output using pure NumPy. First, let's recap your example to make sure we're on the same page:

from sklearn.neighbors import kneighbors_graph
X = [[0, 1], [3, 4], [7, 8]]
A = kneighbors_graph(X, 2, mode='distance', include_self=True)
A.toarray()

Output:
array([[0. , 4.24264069, 0. ],
[4.24264069, 0. , 0. ],
[0. , 5.65685425, 0. ]])

What kneighbors_graph Does Under the Hood #

Let's unpack the parameters first:

• X: Your input sample matrix (each row is a data point)
• n_neighbors=2: Each sample needs its 2 closest neighbors
• mode='distance': The adjacency matrix stores actual distance values instead of binary 0/1 indicators
• include_self=True: Each sample counts itself as one of its own neighbors

The core logic has 3 key steps:

1. Calculate the Euclidean distance between every pair of samples
2. For each sample, select its n_neighbors closest points (including itself)
3. Build an adjacency matrix where only the selected neighbor distances are kept (others are 0)

Replicating with Pure NumPy #

Let's implement each step one by one.

Step 1: Compute Pairwise Euclidean Distances #

We can use NumPy's broadcasting to efficiently calculate all pairwise distances:

import numpy as np

# Convert input to a NumPy array
X = np.array([[0, 1], [3, 4], [7, 8]])

# Calculate pairwise differences (broadcasting handles the matrix expansion)
diff = X[:, np.newaxis, :] - X[np.newaxis, :, :]
# Square the differences, sum across features, then take the square root
distances = np.sqrt(np.sum(diff ** 2, axis=2))

This gives us a 3x3 distance matrix where distances[i,j] is the distance between sample i and sample j.

Step 2: Select the k Closest Neighbors #

We need to find the indices of the 2 smallest distances for each row (including the sample itself, which has a distance of 0). Using np.argpartition is more efficient than full sorting since we only need the top k values:

n_neighbors = 2
# Get indices of the n_neighbors smallest distances for each sample
k_neighbor_indices = np.argpartition(distances, n_neighbors, axis=1)[:, :n_neighbors]

Step 3: Build the Adjacency Matrix #

Initialize a zero matrix and fill in the distances for the selected neighbors:

# Create an empty adjacency matrix with the same shape as the distance matrix
adj_matrix = np.zeros_like(distances)

# Fill in the distances for each sample's k neighbors
for i in range(X.shape[0]):
    adj_matrix[i, k_neighbor_indices[i]] = distances[i, k_neighbor_indices[i]]

Check the Result #

Let's print our NumPy-generated matrix:

print(adj_matrix)

You'll get exactly the same output as scikit-learn:

[[0.         4.24264069 0.        ]
 [4.24264069 0.         0.        ]
 [0.         5.65685425 0.        ]]

Quick Explanation of the Output #

• For sample 0 ([0,1]), its closest neighbors are itself (0) and sample 1 ([3,4], distance ≈4.24)
• For sample 1 ([3,4]), its closest neighbors are itself (0) and sample 0 ([0,1], distance ≈4.24)
• For sample 2 ([7,8]), its closest neighbors are itself (0) and sample 1 ([3,4], distance ≈5.66) — sample 0 is farther away, so it's excluded

内容的提问来源于stack exchange，提问作者Lars