如何对含多列0/1值的DataFrame行执行KMeans聚类并获取行索引

我有一个每行元素仅为0和1的DataFrame，需要把所有行聚成5个簇，同时获取每个簇对应的原始行索引（或能按簇做.groupby操作并保留原始索引）。DataFrame结构示例如下：

0   1   2   3   4   5   6   7   8   9   ... 528 529 530 531 532 533 534 535 536 537
0   0   0   0   0   0   0   0   1   1   1   ... 0   1   1   1   0   0   0   1   0   1
1   0   0   0   0   0   0   0   1   1   1   ... 0   1   1   1   0   0   0   1   0   1
2   0   0   0   0   0   0   0   1   1   1   ... 0   1   1   1   0   0   0   1   0   1
3   0   0   0   0   0   0   0   0   0   0   ... 0   0   0   1   0   0   0   0   1   1
4   0   0   0   0   0   0   0   0   0   0   ... 0   0   0   1   0   0   0   0   0   1
... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
137 0   0   0   0   0   0   0   0   0   0   ... 0   0   0   0   0   1   0   0   0   0
138 1   1   0   0   0   0   0   0   0   1   ... 0   0   0   0   0   1   0   0   0   0
139 1   1   1   0   0   0   0   0   0   0   ... 0   0   0   0   0   1   0   0   0   0
140 1   1   0   0   0   0   0   0   0   1   ... 0   0   0   0   0   1   0   0   0   0
141 1   1   1   0   0   0   0   0   0   0   ... 0   0   0   0   0   1   0   0   0   0

我在Stack Overflow上找到一段解决方案代码：

def cluster(X):
    k_means = KMeans(n_clusters=5).fit(X)
    return X.groupby(k_means.labels_)\
            .transform('mean').sum(1)\
            .rank(method='dense').sub(1)\
            .astype(int).to_frame()

运行后得到结果：

0
0   1
1   1
2   1
3   0
4   0
... ...
137 3
138 1
139 3
140 3
141 3

但我完全看不懂这段代码的作用，也不确定得到的结果是不是每行对应的簇编号。

代码逐行拆解 #

这段代码确实是基于KMeans做聚类生成簇编号，但多了一步编号重排操作，具体逻辑如下：

• KMeans聚类：k_means = KMeans(n_clusters=5).fit(X)
  用KMeans把数据分成5个簇，k_means.labels_就是每行对应的原始簇编号（注意KMeans的原始编号是随机分配的0-4，没有固定顺序）。
• 簇内均值广播：X.groupby(k_means.labels_).transform('mean')
  按原始簇分组，计算每个簇的均值向量，再把这个均值向量复制给该簇的每一行——同一簇的所有行都会被替换成簇均值。
• 均值求和：.sum(1)
  对每行的均值向量求和，同一簇的行这个求和结果完全一致，相当于每个簇的“特征均值总和”。
• 重排簇编号：.rank(method='dense').sub(1)
  对这些总和值做稠密排名（相同值排名相同），再减1转换成从0开始的整数——这一步把KMeans随机的原始编号，改成了按“均值总和”从小到大排序的新编号。
• 格式转换：.astype(int).to_frame()
  把结果转成整数，再包装成DataFrame返回。

结果是不是簇编号？ #

是，返回的数值就是每行对应的簇编号，但这个编号不是KMeans生成的原始编号，而是经过排序重排后的：编号越小，对应的簇的均值总和越小（对于0-1数据来说，均值总和其实就是簇内每行1的数量的平均值，简单说就是簇的“密度”）。

更简洁的实现方法 #

如果只是要获取簇编号、按簇分组并保留原始索引，完全不用这么复杂，直接用KMeans的labels_就行：

from sklearn.cluster import KMeans

def cluster_simple(X):
    k_means = KMeans(n_clusters=5).fit(X)
    # 给原数据添加簇编号列，保留原始索引
    X_with_cluster = X.assign(cluster_id=k_means.labels_)
    # 按簇分组，直接操作分组对象就能获取对应簇的所有行
    cluster_groups = X_with_cluster.groupby('cluster_id')
    return X_with_cluster, cluster_groups

用这个函数的好处：

• X_with_cluster保留了所有原始数据和索引，能直观看到每行对应的簇；
• cluster_groups可以直接调用，比如cluster_groups.get_group(2)就能取出编号为2的簇的全部行，原始索引也会保留。

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