如何单次扫描高效拆分Pandas DataFrame为两个子DataFrame？

好问题！当处理大规模数据集时，两次扫描DataFrame确实会带来不必要的性能开销——尤其是当DataFrame占用内存很高的时候，重复遍历会浪费不少资源。下面分享两种真正实现单次扫描完成拆分的方案：

方法1：用groupby按分组键一次性拆分 #

我们可以把布尔条件作为分组依据，groupby内部会在单次遍历中完成数据分组，之后直接提取各组即可：

import pandas as pd

df = pd.DataFrame({'A': [1,2,3,4], 'B': [1,4,9,16]})

# 生成分组标识：True对应B<5，False对应B>=5
group_flag = df['B'] < 5
# 按标识分组，groupby仅需扫描一次数据
split_groups = dict(tuple(df.groupby(group_flag)))

df1 = split_groups[True]  # 筛选出B<5的部分
df2 = split_groups[False] # 筛选出B>=5的部分

这种方式下，groupby会一次遍历完所有数据，把行分配到对应的组里，完全避免了两次索引切片带来的重复遍历开销。

方法2：结合numpy索引实现极致性能 #

如果追求更底层的性能优化，可以借助numpy的数组操作，先一次性定位所有符合条件的索引，再提取数据：

import pandas as pd
import numpy as np

df = pd.DataFrame({'A': [1,2,3,4], 'B': [1,4,9,16]})

# 用numpy数组生成掩码，比pandas Series的运算更快
mask = df['B'].values < 5
# 一次性获取两类数据的索引位置
idx_small = np.where(mask)[0]
idx_large = np.where(~mask)[0]

# 通过iloc提取数据，底层基于预定位的索引，无需重复扫描
df1 = df.iloc[idx_small]
df2 = df.iloc[idx_large]

这里的关键是用np.where一次性找出所有符合条件的索引，后续的iloc只是根据索引提取数据，不会再遍历整个DataFrame，真正做到了单次扫描完成拆分。

和常规方法的核心差异 #

常规的df1 = df[df['B'] <5] + df2=df[df['B']>=5]，会分别触发两次布尔索引操作——每次索引都要完整遍历DataFrame来筛选行。而上面两种方案，不管是groupby的单次分组，还是numpy的一次性索引定位，都只需要遍历DataFrame一次，在数据量越大的场景下，性能提升越显著。

另外你提到的先计算掩码再切片的方法，虽然掩码只生成一次，但df[m]和df[~m]仍然会分别遍历DataFrame来筛选子集，本质上还是两次数据扫描；而上面的方案则从根源上减少了遍历次数。

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