FastPrepCosine K-Means Final Labels (MLE OA)

Cosine K-Means Final Labels (MLE OA)

Pinterest logoPinterest● MediumINTERNOA

Problem statement

Implement K-Means using cosine distance. Given data, a cluster count k, initialCentroids, and iterations, return one final cluster label per point in the original data order. Cluster IDs are 0 through k - 1 and keep the initial centroid order.

For two nonzero vectors, cosine distance is 1 minus their dot product divided by the product of their Euclidean norms. Use distance 0 for two zero vectors, and distance 1 when exactly one vector is zero. Clamp computed distances to the range [0,2] to handle roundoff. For each point, find the global minimum distance, then choose the smallest cluster ID whose distance is at most that minimum plus 1e-12.

Interpret coordinates as IEEE754 binary64 values. For each cluster and coordinate, accumulate assigned original point values in input order using binary64 additions, then divide by the cluster count in binary64. Use double precision for distance calculations as well.

Each round assigns every point using the current centroids, then simultaneously replaces each nonempty cluster centroid with the coordinate-wise arithmetic mean of its assigned original points. Keep the previous centroid for an empty cluster. Do not normalize points or centroids before computing these means. Perform exactly iterations rounds, then assign all points once more using the final centroids and return those labels. When iterations is 0, use the initial centroids without any mean update.

Function

cosineKMeans(data: double[][], k: int, initialCentroids: double[][], iterations: int) → int[]

Examples

Example 1

data = [[1.0,0.0],[2.0,0.0],[0.0,1.0],[0.0,2.0]]k = 2initialCentroids = [[1.0,0.0],[0.0,1.0]]iterations = 1return = [0,0,1,1]

The first two points have cosine distance 0 to centroid 0, and the last two have distance 0 to centroid 1. After the mean update the centroids are [1.5,0] and[0,1.5], so the final labels remain [0,0,1,1].

Example 2

data = [[0.0,0.0],[1.0,0.0],[-1.0,0.0]]k = 2initialCentroids = [[0.0,0.0],[1.0,0.0]]iterations = 0return = [0,1,0]

There are no mean updates. The zero point selects the zero centroid. The positive point selects centroid 1. The negative point is distance 1 from the zero centroid and distance 2 from centroid 1, so it selects centroid 0.

Constraints

  • 1 <= data.length <= 1000.
  • 1 <= data[i].length <= 10 and 1 <= initialCentroids[i].length <= 10.
  • 1 <= k <= 20 and k <= data.length.
  • initialCentroids.length == k.
  • All data rows and centroid rows have the same length.
  • Every coordinate is finite and between -10000 and 10000, inclusive.
  • 0 <= iterations <= 50.

More Pinterest problems

See Pinterest hiring insights
Saved
public int[] cosineKMeans(double[][] data, int k, double[][] initialCentroids, int iterations) {
    // Write your solution here.
}
Checking account…