如何在Pandas DataFrame中按条件弹出指定行（单语句实现）

场景与需求 #

现有如下Pandas DataFrame：

import pandas as pd
df1 = pd.DataFrame({'A':[1,1,1,1,2,2,2,2],'B':[1,2,3,4,5,6,7,8]})
print(df1)
# 输出：
#    A  B
# 0  1  1
# 1  1  2
# 2  1  3
# 3  1  4
# 4  2  5
# 5  2  6
# 6  2  7
# 7  2  8

需要从df1中弹出A == 2的前2行，同时修改原DataFrame，最终得到：

print(df1)
# 输出：
#    A  B
# 0  1  1
# 1  1  2
# 2  1  3
# 3  1  4
# 6  2  7
# 7  2  8

print(df2)
# 输出：
#    A  B
# 4  2  5
# 5  2  6

注：Pandas原生pop()仅支持弹出整列，无法实现行级弹出，需用以下方法替代。

解决方案 #

可以通过筛选目标行索引→复制到新DataFrame→从原DataFrame删除对应索引的逻辑实现，单语句写法如下：

# 先获取目标行索引，再完成赋值与删除操作
idx = df1[df1['A'] == 2].iloc[:2].index; df2, df1 = df1.loc[idx].copy(), df1.drop(idx)

代码解释 #

1. df1[df1['A'] == 2].iloc[:2].index：筛选出A == 2的前2行的索引（此处为[4,5]）
2. df1.loc[idx].copy()：将目标行复制到df2（用copy()避免后续修改原DataFrame时影响df2）
3. df1.drop(idx)：从原DataFrame中删除目标行，得到修改后的df1

验证结果 #

运行上述代码后，打印df1和df2即可得到需求中的结果。

内容的提问来源于stack exchange，提问作者Oblomov