Matlab转Numpy：矩阵与标量差异下的KNN搜索代码迁移问题

Got it, let's translate your MATLAB k-NN search code to NumPy step by step, while preserving the exact behavior you observed in your example.

First, understand the original MATLAB logic #

Your MATLAB script implements a k=1 nearest neighbor search with a specific rule: for the k-th query point (from matrix Q), it sets the distance to the k-th candidate point (from matrix R) to inf (so it never selects itself when Q and R are the same matrix). The core steps are:

1. Calculate squared Euclidean distances between each query point and all candidate points
2. Exclude the self-match (for k-th query ↔ k-th candidate)
3. Find the smallest distance and its 1-based index (matching MATLAB's indexing style)

NumPy Implementation #

We'll use NumPy's vectorized operations to avoid slow Python loops (critical for performance with large matrices):

import numpy as np

def knnsearch(Q, R, K):
    # Convert inputs to NumPy arrays and ensure float type to avoid integer issues
    Q = np.asarray(Q, dtype=np.float64)
    R = np.asarray(R, dtype=np.float64)
    
    N, M = Q.shape  # N queries, M features per point
    L = R.shape[0]  # L candidate points
    
    # Initialize result arrays to match MATLAB's output shape
    idx = np.zeros((N, K), dtype=np.int64)
    D = np.zeros((N, K), dtype=np.float64)
    
    # Vectorized squared distance calculation (no nested loops!)
    # Shape becomes (N, L) where dists[i,j] is the distance between Q[i] and R[j]
    dists = np.sum((Q[:, None, :] - R[None, :, :]) ** 2, axis=-1)
    
    # Replicate MATLAB's d(k)=inf: exclude k-th candidate for k-th query
    # Note: This assumes N == L (like your example where Q and R are same size)
    # If N != L, this line would throw an error (just like the original MATLAB code)
    if N == L:
        np.fill_diagonal(dists, np.inf)
    
    # Handle k=1 case (your example scenario)
    if K == 1:
        # Get minimum distance for each query
        D[:, 0] = np.min(dists, axis=1)
        # Get index of minimum distance (add 1 to match MATLAB's 1-based indexing)
        idx[:, 0] = np.argmin(dists, axis=1) + 1
    else:
        # Optional: Extend for K>1 (find top K nearest neighbors efficiently)
        # Use argpartition to avoid sorting the entire distance matrix
        top_k_indices = np.argpartition(dists, K, axis=1)[:, :K]
        top_k_dists = np.take_along_axis(dists, top_k_indices, axis=1)
        
        # Sort the top K results by distance to get ordered neighbors
        sort_order = np.argsort(top_k_dists, axis=1)
        idx = np.take_along_axis(top_k_indices, sort_order, axis=1) + 1
        D = np.take_along_axis(top_k_dists, sort_order, axis=1)
    
    return idx, D

Test the Code (Matching Your Example) #

Run this to confirm it produces the same output as your MATLAB script:

Q = np.eye(5)
R = np.eye(5)
idx, D = knnsearch(Q, R, 1)

print("idx:", idx.flatten())  # Output: [2 1 1 1 1]
print("D:", D.flatten())      # Output: [2. 2. 2. 2. 2.]

Key Notes #

1. Vectorization: The distance calculation replaces MATLAB's nested loops with a single vectorized operation, which is drastically faster for large datasets.
2. Indexing: NumPy uses 0-based indexing, so we add +1 to the result indices to match MATLAB's 1-based output format.
3. Self-Exclusion: The np.fill_diagonal line replicates the d(k)=inf behavior from MATLAB, which only works when Q and R have the same number of rows (just like the original code).
4. Scalability: The code includes an optional extension for K>1, using np.argpartition to efficiently find the top K nearest neighbors without sorting the entire distance matrix.

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