Как использовать numpy где на каждом элементе тестового массива без цикла for? - PullRequest
0 голосов
/ 15 июня 2019

Я хотел бы использовать функцию numpy в подпрограмме без использования цикла for.Рассмотрим пример, приведенный ниже:

import numpy as np

data = np.linspace(1, 10, 10).astype(int) 
test_elements = np.array([1, 2])
for test_elem in test_elements:
    print(np.where(test_elem == data))

...

(array([0]),)
(array([1]),)

Я читал другие посты и документальную документацию.Общие предложения, кажется, используют np.roll, чтобы бросить test_elements, или используют другие уловки, такие как шаги (которые я не полностью понимаю).Я подумал, что было бы проще использовать np.vectorize для векторизации функции, но у меня есть ощущение, что это излишне для проблемы и что должно быть более простое решение.Любое руководство будет оценено?

1 Ответ

1 голос
/ 15 июня 2019

Вот один метод, использующий .outer и np.split. Я сделал этот пример немного более интересным.

data = np.linspace(1, 5, 10).astype(int) 
test_elements = np.array([1, 2, 4, 6])
y, x = np.where(np.equal.outer(test_elements,data))
np.split(x, y.searchsorted(np.arange(1,test_elements.size)))
# [array([0, 1, 2]), array([3, 4]), array([7, 8]), array([], dtype=int64)]

Некоторые объяснения:

np.equal.outer(test_elements,data)

совпадает с

 test.elements[:,None] == data[None,:]

Итак, это двумерный логический массив, строки которого равны логическим массивам test_elem == data, встречающимся в вашем цикле for.

where возвращает два индексных массива из этого, по одному для каждой координаты. x - это быстро меняющаяся координата, равная значениям, которые возвращает 1d where вашего цикла for, но все они склеены в один длинный вектор. y - это медленная координата изменения, ее значения упорядочены и могут использоваться для группировки / разделения x. searchsorted, возможно, не самый эффективный способ сделать это, но он прост и правильно обрабатывает пустые строки.

...