You are given one-dimensional sorted data points. Write a program to cluster these points into two clusters such that the WCSS (Within-Cluster Squared Sum) metric is minimized. Try to come up with O(n) solution.

Note : Ensure the clustering results exhibit repeatability, meaning that if the program is executed multiple times, all the data points within a cluster will always remain in the same cluster. Additionally, guarantee optimality in the clustering results. These characteristics are not ensured in the standard K-Means algorithm, where cluster assignments can vary with each execution.

Example 1 : [1, 5, 7, 8, 15, 29, 100] will be clustered into
cluster 1 : [1, 5, 7, 8, 15]
cluster 2 : [29, 100]
Inputs

• x : sorted list of data points
• K = 2, divide into 2 clusters

Output

• Two lists
  • cluster 1 : All the data points in cluster 1
  • Cluster 2 : All the data points in cluster 2
• Min value of WCSS

Solution

---

WCSS can be computed in O(1) by pre computing the prefix_sum and prefix squared sum

Python Code:

import numpy as np
class Solution:

    def calculate_wcss(self, points, l, r, prefix_sum, prefix_sq_sum):
        """Calculate WCSS for points[l:r+1] using prefix sums."""
        n = r - l + 1
        if n == 0:
            return 0
        total_sum = prefix_sum[r + 1] - prefix_sum[l]
        total_sq_sum = prefix_sq_sum[r + 1] - prefix_sq_sum[l]
        mean = total_sum / n
        wcss = total_sq_sum - 2 * mean * total_sum + n * (mean ** 2)
        return wcss

    def two_cluster_partition(self, points):
        """Find the optimal partition for two clusters."""
        n = len(points)

        # Precompute prefix sums for efficient WCSS calculation
        prefix_sum = np.zeros(n + 1)
        prefix_sq_sum = np.zeros(n + 1)
        for i in range(n):
            prefix_sum[i + 1] = prefix_sum[i] + points[i]
            prefix_sq_sum[i + 1] = prefix_sq_sum[i] + points[i] ** 2

        # Iterate over possible partitions and compute WCSS
        min_wcss = float('inf')
        best_partition = -1
        for p in range(n - 1):
            wcss_left = self.calculate_wcss(points, 0, p, prefix_sum, prefix_sq_sum)
            wcss_right = self.calculate_wcss(points, p + 1, n - 1, prefix_sum, prefix_sq_sum)
            total_wcss = wcss_left + wcss_right

            if total_wcss < min_wcss:
                min_wcss = total_wcss
                best_partition = p

        # Create clusters based on the best partition
        cluster1 = points[:best_partition + 1]
        cluster2 = points[best_partition + 1:]
        return min_wcss, [cluster1, cluster2]