Как вернуть значения столбцов с помощью ввода других значений столбцов той же строки, используя панд?

Question

У меня есть такой фрейм данных:

df
col1     col2     col3     col4
 1         2        P        Q
 4         2        R        S
 5         3        P        R

Я хочу создать функцию, которая возвращает значения col1 и col2 с помощью ввода значений col3 и col4,

например, если функция f, вывод f ([P, Q]) будет выглядеть следующим образом:

col1    col2
 1       2

Как это сделать наиболее эффективно, используя панд?

Answer 1

Просто используйте логическое маскирование:

def f(first, second):
    return df.loc[(df['col3'] == first) & (df['col4'] == second), ['col1', 'col2']]

Answer 2

Если нужен самый эффективный способ, сравните numy массивы:

def f(a, b):
    #pandas 0.24+ 
    mask = (df['col3'].to_numpy() == a) & (df['col4'].to_numpy() == b)
    #all pandas versions yet
    #mask = (df['col3'].values == a) & (df['col4'].values == b)
    return  df.loc[mask, ['col1','col2']]

Производительность : Зависит от данных, количества строк, количества совпадающих строк, но, как правило, здесь сравнение 1d-массивов выполняется быстрее:

np.random.seed(123)
N = 10000
L = list('PQRSTU')
df = pd.DataFrame({'col1': np.random.randint(10, size=N),
                   'col2': np.random.randint(10, size=N),
                   'col3': np.random.choice(L, N),
                   'col4': np.random.choice(L, N)})
print (df)

def f(a, b):
    #pandas 0.24+ 
    mask = (df['col3'].to_numpy() == a) & (df['col4'].to_numpy() == b)
    #all pandas versions yet
    #mask = (df['col3'].values == a) & (df['col4'].values == b)
    return  df.loc[mask, ['col1','col2']]

def f1(first, second):
    return df.loc[(df['col3'] == first) & (df['col4'] == second), ['col1', 'col2']]

In [91]: %timeit (f('P', 'Q'))
2.05 ms ± 13.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [92]: %timeit (f1('P', 'Q'))
3.52 ms ± 24.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 3

**Simple line of code can do this**

Вместо 'P' и 'Q' вы должны поместить значения, которым хотите соответствовать.

df[(df.col3 == 'P') & (df.col4 == 'Q')][col1,col2]

Answer 4

Вы можете попробовать код ниже:

def func(x):
    series = f(x['col3'], c['col4'])
    return series.append(x)

dataframe = dataframe.apply(lambda x: func(x))