python - 在一次扫描中将 DataFrame 拆分为 2

Question

假设有以下数据框:

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

如果我必须将其拆分为 2 个数据帧，我们可以按照以下方式进行

df1 = df[df['B'] < 5]
df2 = df[df['B'] >= 5]

但是这里 df 会被扫描两次。有没有办法在单个有效行中将 DataFrame 分成 2 个；一次扫描？

编辑:甚至@jezrael 的建议也有类似的表现:

m = df['B'] < 5
#better performance with comparing numpy array
#m = df['B'].values < 5
df1 = df[m]
df2 = df[~m]  

Answer 1

您可以使用 groupby 如下所示:

df = pd.DataFrame({'A': [1,2,3,4], 'B': [1,4,9,16]})
sampled_dfs = [x for _, x in df.groupby(df['B']<5)]
print(sampled_dfs[0])
print(sampled_df[1])

输出:

  A B
2 3 9
3 4 16

  A B
0 1 1
1 2 4