基于PySpark实现含缺失值的Partial K-means聚类（SPSS PAIRWISE等效）

一、先过滤缺失值数量超过4的样本 #

先给每行统计目标列的缺失数，再筛选出符合要求的样本。假设你的8个目标列是col1到col8，替换成你的实际列名即可：

from pyspark.sql import functions as F

# 定义要处理的8个目标列
target_cols = ["col1", "col2", "col3", "col4", "col5", "col6", "col7", "col8"]

# 计算每行的缺失值数量
df_with_miss_count = df.withColumn(
    "missing_count",
    sum(F.when(F.col(c).isNull(), 1).otherwise(0) for c in target_cols)
)

# 过滤掉缺失数>4的样本，最后删除计数列
filtered_df = df_with_miss_count.filter(F.col("missing_count") <= 4).drop("missing_count")

二、将分类标签转为数值特征 #

K-means只能处理数值型数据，把strongly agree/agree/disagree转成对应数值（比如3/2/1），保留原有的NaN值：

# 逐个列做数值映射
for col_name in target_cols:
    filtered_df = filtered_df.withColumn(
        f"{col_name}_num",
        F.when(F.col(col_name) == "strongly agree", 3)
          .when(F.col(col_name) == "agree", 2)
          .when(F.col(col_name) == "disagree", 1)
          .otherwise(F.col(col_name))  # 保留NaN不做转换
    )

# 收集所有数值化后的列名，后续聚类用
num_cols = [f"{c}_num" for c in target_cols]

三、实现SPSS风格的Partial K-means（MISSING=PAIRWISE） #

PySpark ML原生的K-means会直接丢弃带缺失值的样本，不支持SPSS的成对缺失处理逻辑。下面是自定义实现的核心步骤，完全对应MISSING=PAIRWISE——计算距离时只取样本和质心都非空的特征，更新质心时只取对应类中非空样本的均值：

import numpy as np
from pyspark.sql.types import ArrayType, DoubleType, IntegerType

# 1. 初始化8个质心：从样本里随机选8个非全空的样本作为初始质心（保证质心有有效数据）
initial_centroids = filtered_df.select(num_cols).na.drop(how="all").sample(False, 0.001, seed=42).limit(8).collect()
centroids = [np.array([row[c] for c in num_cols]) for row in initial_centroids]

# 2. 定义计算成对缺失欧氏距离的UDF：返回样本所属的聚类标记（1-8）
def pairwise_cluster_assign(sample_array, centroids):
    min_dist = float("inf")
    closest_cluster = -1
    for idx, centroid in enumerate(centroids):
        # 找出样本和质心都非空的特征索引
        valid_features = ~np.logical_or(np.isnan(sample_array), np.isnan(centroid))
        if np.sum(valid_features) == 0:
            continue  # 没有共同有效特征，跳过该质心
        # 计算有效特征的欧氏距离
        dist = np.linalg.norm(sample_array[valid_features] - centroid[valid_features])
        if dist < min_dist:
            min_dist = dist
            closest_cluster = idx + 1  # 聚类标记从1开始，对应1-8类
    return closest_cluster

distance_udf = F.udf(lambda x: pairwise_cluster_assign(np.array(x), centroids), IntegerType())

# 3. 迭代执行聚类（最大迭代次数设为10，可根据需求调整）
max_iterations = 10
for i in range(max_iterations):
    # 将数值列转为数组，方便UDF处理
    df_with_features = filtered_df.withColumn("features", F.array(*num_cols).cast(ArrayType(DoubleType())))
    # 给每个样本分配聚类标记
    df_with_clusters = df_with_features.withColumn("cluster_label", distance_udf(F.col("features")))
    
    # 更新质心：每个聚类的质心取对应类中各特征的非空均值
    new_centroids = []
    for cluster_id in range(1, 9):
        cluster_data = df_with_clusters.filter(F.col("cluster_label") == cluster_id)
        centroid_vals = []
        for col in num_cols:
            mean_val = cluster_data.select(F.mean(col)).first()[0]
            centroid_vals.append(mean_val if mean_val is not None else np.nan)
        new_centroids.append(np.array(centroid_vals))
    
    # 检查质心是否收敛（变化小于阈值则提前停止）
    if np.allclose(centroids, new_centroids, rtol=1e-3, atol=1e-3, equal_nan=True):
        print(f"聚类在第{i+1}次迭代收敛")
        break
    centroids = new_centroids

# 最终结果：保留原始列和聚类标记列
final_df = df_with_clusters.drop("features")

关键说明： #

• 初始质心选的是随机非全空样本，你也可以根据业务场景手动指定初始值
• 距离计算用的是欧氏距离，和SPSS K-means默认一致
• 对于完全没有共同有效特征的样本，当前代码返回-1，你可以根据需求改成丢弃或归为特定类

四、SPSS聚类代码对应转换说明 #

假设你的SPSS代码是这样：

CLUSTER
/VARIABLES=col1 col2 ... col8
/METHOD=KMEANS(CLUSTERS=8)
/MISSING=PAIRWISE
/PRINT=INITIAL FINAL
/SAVE CLUSTER.

对应的PySpark逻辑就是上面三步：先过滤缺失过多的样本，数值化分类特征，自定义实现带成对缺失处理的K-means，最后生成聚类标记列。

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