KMeansUpdateCentroids extends ComputeFunction { static int updateCentroids(FastDistanceMatrixData matrixData, int k, int vectorSize, double[] buffer, FastDistance distance) { int index = 0; DenseMatrix matrix = matrixData.getVectors(); double[] data = matrix.getData(); Arrays.fill(data, 0.0); for (int clusterId = 0; clusterId < k; clusterId++) { int startIndex = clusterId * (vectorSize + 1); double weight = buffer[startIndex + vectorSize]; if (weight == 0) { continue; } BLAS.axpy(vectorSize, 1.0 / weight, buffer, startIndex, data, index * vectorSize); index++; } distance.updateLabel(matrixData); return index; } KMeansUpdateCentroids(FastDistance distance); @Override void calc(ComContext context);  }
@Test public void updateCentroidsTest(){ int vectorSize = 2; int len = 5; EuclideanDistance distance = new EuclideanDistance(); List<Row> list = new ArrayList<>(); for (int i = 0; i < len; i++) { Vector vec = DenseVector.ones(vectorSize).scale(i); list.add(Row.of(vec)); } FastDistanceMatrixData matrixData = (FastDistanceMatrixData) distance.prepareMatrixData(list, 0).get(0); double[] buffer = new double[]{0.5, 0.5, 2.0, 2.5, 2.5, 2.0, 2.0, 2.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; int k = updateCentroids(matrixData, len, vectorSize, buffer, distance); Assert.assertEquals(k, 3); double[] expect = new double[]{0.25, 0.25, 1.25, 1.25, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0}; double[] predict = matrixData.getVectors().getData(); for (int i = 0; i < expect.length; i++) { Assert.assertEquals(expect[i], predict[i], 0.01); } }