如何根据数据框中的列值获取特定行数

Question

假设我有这种方式的MNIST数据集。

df = pd.read_csv('data/train.csv') 
data = df.loc[df['label'].isin([1,6])] 

我想只选择那些行其列[“标签”] == 1或6

不过，我想从每列仅获得500行[“标签”]

我该怎么做？

Answer 1

使用GROUPBY第一，然后过滤器，即

ndf= df.groupby('label').head(500) 
data = ndf.loc[ndf['label'].isin([1,6])] 

Answer 2

您可以将它们分组，然后选择要为每个值数：

data = df.loc[df['label'].isin([1,6])].groupby('label').head(500)