Compare and store elements of multidimensional array to two new arrays

Viewed 28

Assume I have the following simple array:

my_array = np.array([[1,2],[2,4],[3,6],[2,1]])

which corresponds to another parent array:

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

Of course, there is a function that maps parent_array to np.array but it is not important what function this is.

Goal: I want to use this my_array so as to create two new arrays A and B by iterating each row of my_array: for row i if the value of the first column of my_array[i] is greater than the value of the second column I will store parent_array[i] in A . Otherwise I will store parent_array[i] in B (if the value of the second column in my_array[i] if bigger).

So for the case above the result would be:

A = [3]

because only in the 4-th value of my_array the first column has greater value and

  B = [0,1,2]

because the in the first three rows the second column has greater value.

Now, although I know how to save the greater element in a row of columns to a new array, the fact that each row in my_array is associated with a row in parent_array is confusing me. I don't know how to correlate them.


Summary: I need therefore to associate each row of parent_array to each row of my_array and then if check row by row the latter and if the value of the first column is greater in my_array[i] I save parent_row[i] in A while if the second column is greater in my_array[i] I save parent_row[i] in B.

1 Answers

Use boolean array indexing for this: create boolean condition array by comparing values from 1st and 2nd column of my_array and then use it to select values from parent_array:

cond = my_array[:,0] > my_array[:,1]
A, B = parent_array[cond], parent_array[~cond]
A
# [3]
B
# [0 1 2]
Related