从`sklearn`集群标签输出制作集群列表

如何解决从`sklearn`集群标签输出制作集群列表

sklearn 的聚类输出生成标签列表（其中第 i 个元素将用它所在的聚类进行标记。现在，假设我想要一个聚类列表。有一个相当这样做的明显方法：

def clustarrays(labs):
    howmany = np.max(labs)+1
    results = [[] for i in range(howmany)]
    for i in range(len(labs)):
        cnum = labs[i]
        results[cnum].append(i)
    return results


ll = [1,2,3,5,5]

clustarrays(ll)
[[3,4],[0],[1],[2],[],[5,6,7]]

这一切都很好，但对于大型数据集来说这将非常慢。有没有更以 numpy 为中心的方法来做到这一点？

解决方法

如果你想要纯 numpy 使用：

def clustarrays(labs):
    return np.split(np.argsort(labs),np.unique(labs,return_counts=True)[1].cumsum())

输出：

[array([3,4]),array([0]),array([1]),array([2]),array([5,6,7]),array([],dtype=int64)]

我建议熊猫：

import pandas as pd
def clustarrays(labs):
  df = pd.DataFrame({'labs':labs})
  return df.groupby(df.labs).groups

输出：

{0: Int64Index([3,4],dtype='int64'),1: Int64Index([0],2: Int64Index([1],3: Int64Index([2],5: Int64Index([5,7],dtype='int64')}