python - 如何根据特定对从数据框中选择结果而不进行迭代？

Question

我想根据一些特定的对从数据帧中查询(或定位)子数据帧。

使用迭代来做到这一点很容易，但速度很慢。

import pandas as pd
df=pd.DataFrame([[1,2,3], [1,5,6], [7,8,9], [2,3,8]], columns=['x','y','z'])
df
Out[4]: 
   x  y  z
0  1  2  3
1  1  5  6
2  7  8  9
3  2  3  8

我想获得一个子数据框，其中 (x,y)=(1,2) 和 (x,y)=(1,5) 和 (x,y)=(2,3)，如下所示

Out[5]: 
   x  y  z
0  1  2  3
1  1  5  6
3  2  3  8

我的方法是使用迭代来获取索引:

xy_list=[(1,2),(1,5),(2,3)]
index_list=[]
for x,y in xy_list:
    index_list+=df.query('x==@x & y==@y').index.tolist()
df_sub=df.loc[index_list]
df_sub
Out[6]: 
   x  y  z
0  1  2  3
1  1  5  6
3  2  3  8

有没有什么方法可以在不使用迭代的情况下做到这一点？

Answer 1

您已经很接近了，但您不需要迭代调用query。只需使用 str.join 构建查询字符串，然后进行一次 query 调用即可。

data = [(1, 2), (1, 5), (2, 3)]

pattern = '(' + ') | ('.join(f"x == {a} & y == {b}" for a, b in data) + ')'
pattern
# '(x == 1 & y == 2) | (x == 1 & y == 5) | (x == 2 & y == 3)'

df.query(pattern)

   x  y  z
0  1  2  3
1  1  5  6
3  2  3  8

---

另一个选项是使用 Index.isin 和一些过滤:

df[df.set_index(['x', 'y']).index.isin(data)]

   x  y  z
0  1  2  3
1  1  5  6
3  2  3  8

或者，使用 MultiIndex.from_arrays 构造 MultiIndex:

df[pd.MultiIndex.from_arrays([df['x'], df['y']]).isin(data)]

   x  y  z
0  1  2  3
1  1  5  6
3  2  3  8

结果相同，效率更高。