数据框分类和排序优化问题

Question

我想取dataframe列'category'中的两个字母和另一列'data1'中最大的4个，按照一定的规则排序。我使用了根据类别中的字母进行切片排序的方法，然后使用concat进行合并。代码如下。还有比这更简单的方法吗？

import numpy as np
from time import time
import pandas as pd

df = pd.DataFrame()
n = 200
df['category'] = np.random.choice(('A', 'B'), n)
df['data1'] = np.random.randint(1, 10000, len(df))
df['data2'] = np.random.randint(1, 10000, len(df))
a = df[df['category'] == 'A'].sort_values(by='data1', ascending=False).head(4)
b = df[df['category'] == 'B'].sort_values(by='data1', ascending=False).head(4)
df = pd.concat([a, b]).sort_values(by=['category', 'data1'], ascending=[True, False]).reset_index(drop=True)
print(df)

  category  data1  data2
0        A   9963   7174
1        A   9921   1787
2        A   9870   6105
3        A   9779   7189
4        B   9938   4259
5        B   9714   9362
6        B   9622   7033
7        B   9600   5718

我又试了lambda，代码如下：

import numpy as np
from time import time
import pandas as pd


df = pd.DataFrame()
n = 200
df['category'] = np.random.choice(('A', 'B'), n)
df['data1'] = np.random.randint(1, 10000, len(df))
df['data2'] = np.random.randint(1, 10000, len(df))
df = df.groupby('category').apply(lambda x: x[x['data1'].isin(x['data1'].nlargest(4))]).reset_index(drop=True)
df = df.sort_values(by=['category', 'data1'], ascending=[True, False]).reset_index(drop=True)
print(df)

Answer 1

用途：

#for same random data for compare
np.random.seed(2021)


df = pd.DataFrame()
n = 200
df['category'] = np.random.choice(('A', 'B'), n)
df['data1'] = np.random.randint(1, 10000, len(df))
df['data2'] = np.random.randint(1, 10000, len(df))
a = df[df['category'] == 'A'].sort_values(by='data1', ascending=False).head(4)
b = df[df['category'] == 'B'].sort_values(by='data1', ascending=False).head(4)
df1 = pd.concat([a, b]).sort_values(by=['category', 'data1'], ascending=[True, False]).reset_index(drop=True)
print(df1)
  category  data1  data2
0        A   9882   9868
1        A   9855   6701
2        A   9798   1058
3        A   9669   7334
4        B   9973   3668
5        B   9900   4340
6        B   9846   7885
7        B   9659   4933

两列先使用DataFrame.sort_values，再添加GroupBy.head：

df1 = (df.sort_values(by=['category', 'data1'], ascending=[True, False])
         .groupby('category')
         .head(4)
         .reset_index(drop=True))
print (df1)
  category  data1  data2
0        A   9882   9868
1        A   9855   6701
2        A   9798   1058
3        A   9669   7334
4        B   9973   3668
5        B   9900   4340
6        B   9846   7885
7        B   9659   4933