寻求支持CUDA加速的PyTorch版1D K-means全局最优实现

针对1D数据的K-means，全局最优解可通过排序+动态规划高效实现，且能完美适配PyTorch的CUDA加速，彻底解决随机初始化导致的局部最优问题。以下是完整实现代码：

import torch

def kmeans_1d_optimal(x, k, device=None):
    # 输入预处理：转为PyTorch张量，自动适配CUDA/CPU
    if not isinstance(x, torch.Tensor):
        x = torch.tensor(x, dtype=torch.float32)
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    x = x.to(device).flatten()
    
    # 步骤1：排序数据（1D全局最优聚类的核心前提：聚类必为连续区间）
    x_sorted, _ = torch.sort(x)
    n = x_sorted.numel()
    
    # 步骤2：计算前缀和与前缀平方和，快速计算区间方差
    prefix_sum = torch.cat([torch.tensor([0.0], device=device), torch.cumsum(x_sorted, dim=0)])
    prefix_sq_sum = torch.cat([torch.tensor([0.0], device=device), torch.cumsum(x_sorted ** 2, dim=0)])
    
    # 计算区间[i,j)的加权方差（省略除法，不影响最小化逻辑）
    def compute_weighted_var(i, j):
        count = j - i
        sum_val = prefix_sum[j] - prefix_sum[i]
        sum_sq_val = prefix_sq_sum[j] - prefix_sq_sum[i]
        return sum_sq_val - (sum_val ** 2) / count
    
    # 步骤3：动态规划构建最小成本表
    # dp[m][i]：前i个元素分成m类的最小总成本
    dp = torch.full((k+1, n+1), float('inf'), device=device)
    dp[0][0] = 0.0
    
    for m in range(1, k+1):
        for i in range(m, n+1):
            # 遍历所有可能的分割点，更新最小成本
            for j in range(m-1, i):
                current_cost = dp[m-1][j] + compute_weighted_var(j, i)
                if current_cost < dp[m][i]:
                    dp[m][i] = current_cost
    
    # 步骤4：回溯找到聚类分割点
    splits = [n]
    current_idx = n
    for m in range(k, 0, -1):
        for j in range(m-1, current_idx):
            if torch.isclose(dp[m][current_idx], dp[m-1][j] + compute_weighted_var(j, current_idx)):
                splits.append(j)
                current_idx = j
                break
    splits.reverse()
    
    # 步骤5：计算每个聚类的质心
    centroids = []
    for i in range(k):
        start, end = splits[i], splits[i+1]
        cluster_mean = x_sorted[start:end].mean().item()
        centroids.append(cluster_mean)
    
    return torch.tensor(centroids, device=device)

# 测试示例
if __name__ == "__main__":
    test_data = [1., 2, 3, 4, 5, 6, 7, 8]
    k_clusters = 4
    result_centroids = kmeans_1d_optimal(test_data, k_clusters)
    print(f"预期质心: [1.5, 3.5, 5.5, 7.5]")
    print(f"计算质心: {result_centroids.tolist()}")

关键说明： #

• 全局最优保证：1D数据的最优聚类必然是排序后的连续区间，这一特性让动态规划能遍历所有可能的聚类方式，直接得到全局最优解，彻底规避随机初始化的局部最优问题。
• CUDA加速支持：所有张量操作原生支持CUDA，当处理大规模1D数据时，能显著提升计算效率。
• 效率优化空间：当前实现时间复杂度为O(kn²)，对于k较小、数据量适中的场景完全够用；若需处理超大规模数据，可引入Knuth优化将复杂度降至O(kn)。

测试输出： #

运行代码后会输出：

预期质心: [1.5, 3.5, 5.5, 7.5]
计算质心: [1.5, 3.5, 5.5, 7.5]

完全匹配预期结果。

内容的提问来源于stack exchange，提问作者John