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和isin

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

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

尝试一下：

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

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

答案比较

3 个答案: