如何在没有for循环的情况下在测试数组的每个元素上使用numpy？

Question

I would like to use a numpy function in a routine without using a for-loop. Consider the example below:

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]),)

I've read other posts and the numpy documentation. Common suggestions seem to be using np.roll to roll test_elements , or using other tricks such as strides (which I do not completely understand). I thought it might be easier to use np.vectorize to vectorize the function, but I have a feeling that this is overkill for a problem and that there must be a simpler solution. Any guidance would be appreciated?

Answer 1

Here is one method using .outer and np.split I've made the example a bit more interesting.

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)]

Some explanation:

np.equal.outer(test_elements,data)

is the same as

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

So this is a 2d boolean array whose rows equal the boolean arrays test_elem == data occuring in your for loop.

where returns two index arrays from this, one for each coordinate. x is the fast changing coordinate and is equal to the values the 1d where of your for loop returns, but all glued together into one long vector. y is the slow chaning coordinate, it's values are ordered and can be used to group / split up x . searchsorted may not be the most efficient way of doing that but it is simple and handles empty rows correctly.