PDS/homework_1/matlab/knnsearch2.m

function [idx, dst] = knnsearch2(C, Q, k)
% C: Is a mxd matrix (Corpus)
% Q: Is a nxd matrix (Query)
% k: The number of nearest neighbors needded

    % Calculate the distance matrix between C and Q
    % D is an m x n matrix where each element D(i, j) is the distance 
    % between the i-th point in C and the j-th point in Q.
    % k is the number of nearest neighbors to find.
    D = dist2(C, Q);

    % Find the k-nearest neighbors for each query point in Q
    % [~,n] = size(D);
    % for j = 1:n        
    %    [dst(:, j), idx(:, j)] = mink(D(:, j), k);
    % end
    [dst, idx] = mink(D, k, 1);  % mink along dimension 1 for each query point
    
    % Transpose the output to match the knnsearch format
    idx = idx';  % Make idx an n x k matrix
    dst = dst';  % Make dst an n x k matrix

end