python - 选择 pandas DataFrame 中的列子集与元组匹配的位置

Question

这是一个简单的问题，我似乎找不到优雅的解决方案。我正在尝试从单独的列表中选择数据框的行，其中两列形成一对。

例如:

import pandas as pd

df = pd.DataFrame({'a': range(8), 'b': range(8), 'c': list('zyxwvuts')})
pairs = [(4, 4), (5, 6), (6, 6), (7, 9)]

# The data has an arbitrary number of columns, but I just want
# to match 'a' and 'b'
df
    a   b   c
0   0   0   z
1   1   1   y
2   2   2   x
3   3   3   w
4   4   4   v
5   5   5   u
6   6   6   t
7   7   7   s

在这个例子中，我的列表 pairs 在第 4 行和第 6 行包含 df.a 和 df.b 的组合。想要一个干净的方法来获取 df.iloc[[4, 6], :] 给出的数据框。

是否有 pandas 或 numpy 方法可以在不显式循环 pairs 的情况下执行此操作？

---

答案比较

使用广播的解决方案既干净又快速，而且扩展性非常好。

def with_set_index(df, pairs):
    return df.set_index(['a','b']).loc[pairs].dropna()

def with_tuple_isin(df, pairs):
    return df[df[['a','b']].apply(tuple,1).isin(pairs)]

def with_array_views(df, pairs):
    def view1D(a, b): # a, b are arrays
        a = np.ascontiguousarray(a)
        b = np.ascontiguousarray(b)
        void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
        return a.view(void_dt).ravel(), b.view(void_dt).ravel()

    A, B = view1D(df[['a','b']].values, np.asarray(pairs))
    return df[np.isin(A, B)]

def with_broadcasting(df, pairs):
    return df[(df[['a','b']].values[:,None] == pairs).all(2).any(1)]

%timeit with_set_index(df, pairs)
# 7.35 ms ± 119 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit with_tuple_isin(df, pairs)
# 1.89 ms ± 24.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit with_array_views(df, pairs)
# 917 µs ± 17.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit with_broadcasting(df, pairs)
# 879 µs ± 8.85 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 1

tuple with isin

df[df[['a','b']].apply(tuple,1).isin(pairs)]
Out[686]: 
   a  b  c
4  4  4  v
6  6  6  t