将通过numpy数组分配的群集转换为单独的数据集

Question

我有一个非常大的数据集，我正在运行一个聚类模型。聚类输出一个格式如下的numpy数组：

    [ 0 1 2 1 1 0 0 0 1 2 1 0 2 0 1 2 1 0 2 2 0 0 1 ... ]

我想获取原始数据集，并根据数组创建三个数据集。我该怎么做呢？

初始数据集工作：

 import pandas as pd
 pd.options.mode.chained_assignment = None
 raw_data = pd.read_csv("LendingClub2012to2013.csv", low_memory = False, skiprows=[0])

 //Some cleaning done, target leakage removed, dummies created, imputation, etc.

 clean_data = raw_data.drop(text2d + leakage2d + noinfo2d + irr2d, axis = 1)

Answer 1

我假设您的数据集是一个numpy数组。尝试创建蒙版以从原始数据集中选择所需的元素。一些冗长的代码：

# Your original data set (2d numpy array)
orig_data = ...
# The cluster assignments output by the algorithm (1d numpy array)
cluster_assignments = ...

clusters = []
for cluster_id in xrange(3):
    mask = (cluster_assignments == cluster_id)
    clusters.append(orig_data[mask])

更简洁的版本：

clusters = [orig_data[cluster_assignments == id] for id in xrange(3))]

如果你的数据集是pandas Dataframe而不是numpy数组，那么简单地用orig_data[...]替换orig_data.loc[...]。

此代码的输出是一个列表clusters，其中每个元素都是一个数据集，其中只包含其中一个集群的数据。