Для всех людей, которые качаются на векторизованных петлях: у меня есть два NumPy массива формы (N,)
, которые содержат индексы друг к другу. Скажем, у нас есть a = np.asarray([0, 1, 2])
и b = np.array([1, 2, np.nan])
. Функция должна сначала посмотреть на a[0]
, чтобы получить 0
, затем сделать b[0]
, чтобы получить 1
, затем снова a[1]
, чтобы получить 2
, и так далее, пока мы не получим np.nan
. Так что функция просто a[b[a[b[a[0]]]]] = np.nan
. Вывод должен содержать два списка значений, которые были вызваны для a
и b
соответственно. Индексы в b
всегда больше, чем в a
, так что процесс не может застрять.
Я написал простую функцию, которая может сделать именно это (завернутый в numba - 18,2 мкс):
a = np.array([0, 1, 2, 3, 4])
b = np.array([ 2., 3., 4., nan, nan])
lst = []
while True:
if len(lst) > 0:
idx = lst[-1]
else:
idx = 0
if len(lst) % 2 == 0:
if idx < len(a) - 1:
next_idx = a[idx]
if np.isnan(next_idx):
break
lst.append(int(next_idx))
else:
break
else:
if idx < len(b) - 1:
next_idx = b[idx]
if np.isnan(next_idx):
break
lst.append(int(next_idx))
else:
break
Первый список lst[::2]
:
[0, 2]
Второй: lst[1::2]
:
[2, 4]
Есть ли способ векторизовать это? Оба массива на входах, а также оба списка на выходе всегда имеют одинаковую форму.