Calculation of the Distance Matrix in the K-Means Algorithm in MATLAB

Question

Purpose of the code :
To assign the corresponding label of the centroids to the points which are close to it. Below is a graphical (2D) example.

---

Variable X is a matrix, rows represent the points, columns collectively represent the coordinates

idx is the list of centroids assigned to every example

Inefficient MATLAB code to find and assign nearest Centroids

sd2=size(X,2);
t2=[];
for a=1:K
    t1=[];
    for b=1:sd2
        temp=(X(:,b)-centroids(a,b)).^2;
        t1=[t1 temp];
    end
    dmat=sum(t1,2);
    t2=[t2 dmat];
end
[~,idx]=min(t2,[],2);

supposedly faster MATLAB code which worked on a small dataset:

[~,idx]=min(sum(((repmat(X,1,1,K)-reshape(centroids,1,n,K)).^2),2),[],3);

My Question:
Is my second implementation correct? If yes, why doesn't it work on a 300 x 2 matrix X , in the sense that the final array idx produced is not correct. If no, where is the problem (because I tested it with a small sample code just to be sure that I am finding (x-x₁)²+ (y-y₁)² correctly).

speculation:
MATLAB is erroneous because of a multidimensional matrix.

Answer 1

Your Distance Matrix Calculation

You wrote the code (I summarize it for the Distance Matrix):

K = 3;
n = size(mX, 1);
d = size(mX, 2); %<! Number of dimensions per sample
mC = mX(randperm(n, K), :);

% mD = repmat(mX, 1, 1, K) - reshape(mC, 1, n, K); %<! Your code, Will fail because of reshape
mD = sum((reshape(mX, n, 1, d) - reshape(mC.', 1, K, d)) .^ 2, 3); %<! Correctly with no `repmat()`

The reason it won't work is the reshaping operation you're doing.
The number of elements in mC is K x d while your reshape() should keep the number of elements.

In the following I gave you a code snippet how to calculate the matrix mD efficiently.

K-Means

To calculate the distance you shouldn't use repmat() which will allocate new memory. To calculate the Distance Matrix with the 3rd dimension and broadcasting you should do something like:

mD = sum((reshape(mA, numVarA, 1, varDim) - reshape(mB.', 1, numVarB, varDim)) .^ 2, 3);

But a faster way would be:

mD = sum(mA .^ 2).' - (2 * mA.' * mB) + sum(mB .^ 2);

Which utilizes the fast the Matrix Multiplication is probably the most optimized algorithm and that:

$$ {\left\| \boldsymbol{x} - \boldsymbol{y} \right\|}_{2}^{2} = \boldsymbol{x}^{T} \boldsymbol{x} - 2 \boldsymbol{x}^{T} \boldsymbol{y} + \boldsymbol{y}^{T} \boldsymbol{y} $$

For you to have a reference of K-Means algorithm with support for any Distance Function I created a function and ran it on your data:

It will converge very fast for that data:

I also created a small benchmark for various algorithm to compute the Euclidean Squared Distance Matrix.
The first test was between 3 approaches to calculate the distance between a single set of vectors (For example given by a single matrix):

The second test I compared 2 Row / Column variations of the algorithm above:

The Row / Column variations means I test the algorithm in 2 cases:

1. Row Variation - Each data sample is a row in the data matrix (Like your data).
2. Column Variation - Each data sample is a column in the data matrix (In MATLAB it means the sample is contiguous in memory).

As can be seen, it is always better to have data in the column variation.

Remark: Pay attention that the Squared Euclidean Distance doesn't obey the Triangle Inequality which is a requirement for a Metric. It is used for computational efficiency as it doesn't change the K-Means algorithm results compared to Euclidean Distance (The L₂ Norm based distance).

The MATLAB Code is accessible in my StackExchange Code Review Q254186 GitHub Repository.

Remark: The K-Means algorithm, as a non convex optimization, is sensitive to the initialization. In the above code I used random initialization - randomly choosing samples form data to be the initial centroids. There are better approaches out there (See K-Means++ as an example for a more robust initialization).