Python 数据帧分类和排序优化问题

我想在数据框列

'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)

使用：

先由两列使用，然后添加：

---

你的代码非常简洁，你能帮我看看另一个类似的问题吗？

#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

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