Note that this question is not about multiple conditions within a single np.where(), see this thread for that.

I have a numpy array arr1 with some numbers (without a particular structure):

arr0 = \
np.array([[0,3,0],
          [1,3,2],
          [1,2,0]])

and a list of all the entries in this array:

entries = [0,1,2,3]

I also have another array, arr1:

arr1 = \
np.array([[4,5,6],
          [6,2,4],
          [3,7,9]])

I would like to perform some function on multiple subsets of elements of arr1. A subset consts of numbers which are at the same position as arr0 entries with a cetrain value. Let this function be finding the max value. Performing the function on each subset via a list comprehension:

res = [np.where(arr0==index,arr1,0).max() for index in entries]

res is [9, 6, 7, 5]

As expected: 0 in arr0 is on the top left, top right, bottom right corner, and the biggest number from the top left, top right, bottom right entries of arr1 (ie 4, 6, 9) is 9. Rest follow with a similar logic.

How can I achieve this without iteration?

My actual arrays are much bigger than these examples.

CodePudding user response：

With broadcasting

res = np.where(arr0[...,None] == entries, arr1[...,None], 0).max(axis=(0, 1))

The result of np.where(...) is a (3, 3, 4) array, where slicing [...,0] would give you the same 3x3 array you get by manually doing the np.where with just entries[0], etc. Then taking the max of each 3x3 subarray leaves you with the desired result.

CodePudding user response：

It's more convenient to work in 1D case. You need to sort your arr0 then find starting indices for every group and use np.maximum.reduceat.

arr0_1D = np.array([[0,3,0],[1,3,2],[1,2,0]]).ravel()
arr1_1D = np.array([[4,5,6],[6,2,4],[3,7,9]]).ravel()
arg_idx = np.argsort(arr0_1D)
>>> arr0_1D[arg_idx]
array([0, 0, 0, 1, 1, 2, 2, 3, 3])
u, idx = np.unique(arr0_1D[arg_idx], return_index=True)
>>> idx
array([0, 3, 5, 7], dtype=int64)
>>> np.maximum.reduceat(arr1_1D[arg_idx], idx)
array([9, 6, 7, 5], dtype=int32)