Выберите, где подмножество столбцов в панде 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 содержит комбинацию df.a и df.b в строках 4 и 6. Я хотел бы получить простой способ получить фрейм данных, заданный 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

Векторизованный, основанный на представлениях массива -

# 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 2

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 3

Попробуйте это:

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