Как выбрать первую строку, удовлетворяющую условию?
Вот набор данных:
import pandas as pd
test = pd.DataFrame({"par1": [1, 2, 3, 4, 5],
"par2": [11, 12, 13, 14, 15],
"par3": [21, 22, 23, 24, 25],
})
test
Out[45]:
par1 par2 par3
0 1 11 21
1 2 12 22
2 3 13 23
3 4 14 24
4 5 15 25
Как получить первую строку, удовлетворяющую набору условий?
Конкретно у меня есть набор чисел p1, p2, p3. Как выбрать первую строку, для которой верно условие par1>p1 & par2>p2 & par3>p3? Например найти коробку, в которую влезет объект с размерами p1, p2, p3.
Вот мое решение:
test[(test["par1"]>2) & (test["par2"]>13) & (test["par3"]>1)].iloc[0, :]
Out[43]:
par1 4
par2 14
par3 24
Name: 3, dtype: int64
Есть подозрение, что в Pandas есть более эффективное решение . Это как выход из цикла по условию с помощью break. Выбор ВСЕХ записей, удовлетворяющих условию явно избыточен.
Ответы (2 шт):
"более эффективное" решение не будет векторизированным, т.к. придётся использовать цикл. Поэтому может оказаться медленнее чем "менее эффективное". ;)
В качестве альтернативы можно воспользоваться Numba:
from numba import prange, njit, jit, i8
@njit('i8(i8[:, :], i8, i8, i8)')
def fun(arr, p1, p2, p3):
for i in range(len(arr)):
if arr[i, 0] > p1 and arr[i, 1] > p2 and arr[i, 2] > p3:
return i
return -1
тесты:
In [59]: idx = fun(test.to_numpy(), p1, p2, p3)
In [60]: res = test.iloc[idx] if idx >= 0 else None
In [61]: res
Out[61]:
par1 4
par2 14
par3 24
Name: 3, dtype: int64
Для массивов данных более 1.000.000 записей советую воспользоваться GPU (numba + cuda) или cupy.
PS сравнение скорости работы на больших массивах оставлю вам)
Сравнение производительности для фрейма: (5000, 3)
df = pd.concat([test] * 10**3, ignore_index=True)
df.shape
>>> (5000, 3)
Pandas
%timeit df[(df["par1"]>2) & (df["par2"]>13) & (df["par3"]>1)].iloc[0]
>>> 846 µs ± 6.17 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit df.query("par1 > 2 and par2 > 13 and par3 > 1").iloc[0]
>>> 1.87 ms ± 11.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Numba
%%timeit
idx = fun(df.to_numpy(), 2, 13, 1)
res = df.iloc[idx] if idx >= 0 else None
>>> 74.9 µs ± 577 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%%timeit
idx = fun(df.to_numpy(), 2, 13, 1000)
res = df.iloc[idx] if idx >= 0 else None
>>> 8.34 µs ± 63.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
Modin[ray]
%timeit df[(df["par1"]>2) & (df["par2"]>13) & (df["par3"]>1)].iloc[0]
>>> 11.9 s ± 267 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit df.query("par1 > 2 and par2 > 13 and par3 > 1").iloc[0]
>>> 2.3 s ± 645 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Вывод
- Pandas: 1 (baseline)
- Numba: 0.0885 (в 11 раз быстрее Pandas)
- Modin[ray]: 2.7187 (в 2.7187 медленнее Pandas)
PS я ожидал от разрекламированного modin[ray] большего.