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

Question

假设我有这样一个MNIST数据集。

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

我试图只选择列['label'] == 1或6的那些行。

但是，我希望每列只有500行['label']

我该怎么做？

Answer 1

首先使用groupby然后使用filer，即

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)