Как получить максимальное число из n предыдущих чисел вектора A взяв n из вектора B?
Надо производительным способом без циклов.
A = np.array([2, 3, 5, 2, 1, 1, 4, 5, 4, 1, 1, 1], dtype=float)
B = np.array([2, 2, 3, 1, 5, 3, 2, 5, 4, 3, 2, 1])
Ожидаемый результат:
[nan 3. 5. 2. 5. 2. 4. 5. 5. 5. 1. 1.]
Пример решения с циклом:
for i in range(len(A)):
n = B[i]
if i >= n - 1:
chunk = A[(i - n + 1):(i + 1)]
res[i] = np.max(chunk)
else:
res[i] = np.nan
print(res)
Ответы (2 шт):
Я не уверен, что это можно сделать векторизированным способом в принципе.
Но можно задействовать Numba JIT (Just In Time) compiler:
функция с использованием @numba.njit декоратора:
from numba import prange, njit, jit
@njit #(['float64[:](float64[:], int64[:])'])
def numba_max_prevs(arr, prevs):
n = len(arr)
r = np.zeros(n, dtype=np.float64)
r[0] = np.nan
for i in range(1, n):
idx_from = max(i-prevs[i]+1, 0)
r[i] = np.max(arr[idx_from : i+1])
return r
проверка:
arr = A
prevs = B
In [64]: numba_max_prevs(arr, prevs)
Out[64]: array([nan, 3., 5., 2., 5., 2., 4., 5., 5., 5., 1., 1.])
сравнение скорости выполнения:
та же самая функция, но без использования @numba.njit декоратора:
def max_prevs(arr, prevs):
n = len(arr)
r = np.zeros(n, dtype=np.float64)
r[0] = np.nan
for i in range(1, n):
idx_from = max(i-prevs[i]+1, 0)
r[i] = np.max(arr[idx_from : i+1])
return r
замеры времени исполнения для массивов, увеличенных в 1000 раз - 12.000 элементов:
In [65]: arr_big = np.concatenate([arr] * 10**3)
In [66]: prevs_big = np.concatenate([prevs] * 10**3)
In [67]: arr_big.shape
Out[67]: (12000,)
In [68]: %timeit max_prevs(arr_big, prevs_big)
130 ms ± 705 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [69]: %timeit numba_max_prevs(arr_big, prevs_big)
1.07 ms ± 365 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
выигрыш - в 121 раз быстрее:
In [70]: 130 / 1.07
Out[70]: 121.49532710280373
PS я не уверен как должна себя вести функция для первого элемента массива если первым элементов в списке B - будет единица. Думаю с этим вы сами справитесь)
Вообще, тут для улучшения производительности надо менять сам алгоритм. Алгоритм из вопроса имеет квадратичную асимптотику, а задачу можно решить за O(n*lb(n)). Для этого следует воспользоваться деревом отрезков или модицицированным деревом Фенвика.
Лично я писал бы дерево отрезков, но просто потому, что его я писать умею, а с Фенвиком надо разбираться. Никаких причин предпочесть один из этих алгоритмов другому я назвать не могу.