选择 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

试试这个：

df.set_index(['a','b']).loc[pairs].dropna()

Answer 2

基于数组视图的矢量化 -

# https://stackoverflow.com/a/45313353/ @Divakar
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))
out = df[np.isin(A,B)]

给定样本的输出 -

In [263]: out
Out[263]: 
   a  b  c
4  4  4  v
6  6  6  t

如果您正在寻找紧凑/干净的版本，我们还可以利用 broadcasting -

In [269]: df[(df[['a','b']].values[:,None] == pairs).all(2).any(1)]
Out[269]: 
   a  b  c
4  4  4  v
6  6  6  t

Answer 3

tuple 与isin

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