Luca
Luca

Reputation: 10996

python: expanding a numpy array inside function

I have a python function with the following signature:

def merge(segments, indexes):

where segments is a n-d numpy array and indexes is a one dimensional numpy array. Now, I want to call the following function:

np.where((segments == indexes[0]) | (segments == indexes[1]) | 
          ... segments == indexes[n])

However, I am not sure how I can generate this condition dynamically within the where() function call in python.

Upvotes: 1

Views: 100

Answers (1)

Saullo G. P. Castro
Saullo G. P. Castro

Reputation: 58955

Since you have many or conditions, you can use np.in1d() to check if each element of segments exists in indexes:

np.where(np.in1d(segments, indexes).reshape(segments.shape))

Note that the output of in1d() is a flattened array, needing to be reshaped such that where() will return the correct indices.

Upvotes: 1

Related Questions