Есть ли способ векторизовать два массива, указывающих друг на друга в NumPy? - PullRequest
3 голосов
/ 31 января 2020

Для всех людей, которые качаются на векторизованных петлях: у меня есть два 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]

Есть ли способ векторизовать это? Оба массива на входах, а также оба списка на выходе всегда имеют одинаковую форму.

Ответы [ 2 ]

2 голосов
/ 31 января 2020

Это не векторизованное решение, но как решение Numba оно должно быть намного быстрее и проще. Я немного изменил код, чтобы использовать целые числа и -1 вместо np.nan, переключиться на это представление тривиально, например b = np.where(np.isnan(b), -1, b), и это делает код более эффективным. Вместо того, чтобы иметь растущую структуру в функции Numba, я заранее выделяю выходной массив, чтобы l oop мог работать намного быстрее.

import numba as nb

def point_each_other(a, b):
    # Convert inputs to array if they are not already
    a = np.asarray(a)
    b = np.asarray(b)
    # Make output array in advance
    out = np.empty(len(a) + len(b), dtype=a.dtype)
    # Call Numba function
    n = point_each_other_nb(a, b, out)
    # Return relevan part of the output
    return out[:n]

@nb.njit
def point_each_other_nb(a, b, out):
    curr = 0
    i = 0
    while curr >= 0:
        # You can do bad input checking with the following
        # if i >= len(out):
        #     raise ValueError
        # Save current index
        out[i] = curr
        # Get the next index
        curr = a[curr]
        # Swap arrays
        a, b = b, a
        # Advance counter
        i += 1
    # Return number of stored indices
    return i - 1

# Test
a = np.array([0, 1, 2, 3, 4])
b = np.array([2, 3, 4, -1, -1])
out = point_each_other(a, b)
print(out[::2])
# [0 2 4]
print(out[1::2])
# [0 2]
0 голосов
/ 31 января 2020

Не векторизовано, но вот рекурсивное решение:

import numpy as np

a = np.array([0,1,2,3,4])
b = np.array([2,3,4,np.nan, np.nan])

def rec(i,b,a, a_out, b_out):
    if np.isnan(b[i]): return
    else:
        if not np.isnan(b[i]): a_out.append(i)
        rec(int(b[i]), a, b, b_out, a_out)
    return a_out, b_out

print(rec(0,b,a,[],[]))

Выход

([0, 2], [2, 4])
Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...