While this is a neat trick, the claim of "100x" speedup for the centroid computation is somewhat unrealistic, since I doubt that there are many practical applications where one wants to compute 10000 centroids for 10000 data points. The proposed technique is 50 % slower than the naive implementation for a more reasonable k = 10 centroids.
Anyway, JIT compilers are quite good these days. A naive implementation with for loops beats the fastest version in both cases by a factor of 5 (here with numba):
from numba import njit
@njit
def numba_centroids(X_nd, label_n):
n, d = X_nd.shape
k = label_n.max() + 1
c_kd = np.zeros((k, d))
weights = np.zeros(k)
dist_n = np.zeros(n)
for i in range(n):
label = label_n[i]
for j in range(d):
c_kd[label, j] += X_nd[i, j]
weights[label] += 1
c_kd /= weights.reshape(k, 1) + 1e-10
for i in range(n):
label = label_n[i]
squared_difference = 0.0
for j in range(d):
difference = c_kd[label, j] - X_nd[i, j]
squared_difference += difference * difference
dist_n[i] = squared_difference
return c_kd, dist_n