| 42 | } |
| 43 | |
| 44 | array new_means(array data, array clusters, int k) { |
| 45 | int d = data.dims(2); |
| 46 | array means = constant(0, 1, k, d); |
| 47 | array clustersd = tile(clusters, 1, 1, d); |
| 48 | |
| 49 | gfor(seq ii, k) { |
| 50 | means(span, ii, span) = |
| 51 | sum(data * (clustersd == ii)) / (sum(clusters == ii) + 1e-5); |
| 52 | } |
| 53 | |
| 54 | return means; |
| 55 | } |
| 56 | |
| 57 | // kmeans(means, clusters, data, k) |
| 58 | // data: input, 1D or 2D (range > [0-1]) |