zlenyk
zlenyk

Reputation: 1010

Numpy 2d array, select indices satisfying conditions from 2 arrays

I have two 3x3 arrays. One of them indicates if an element is black (let's say 0's - white, 1's black) and another what is the cost of an element. Is there a nice way to get indices of for example all elements that are black and their price is higher than certain value? I know I can use np.where() to select from one array, but how to do it on two (if they have the same shapes)

Upvotes: 0

Views: 2124

Answers (1)

user6655984
user6655984

Reputation:

Following up on the advice of Psidom and rayryeng, I'll add that the output of np.where can be stacked to present a list of indices in the readable "coordinate" notation, as shown below

import numpy as np
a = np.random.randint(0, 2, size=(3,3))
b = np.random.uniform(0, 10, size=(3,3))
print(a)
print(b)
print(np.where(a & (b > 4)))
print(np.vstack(np.where(a & (b > 4))).T)

Random arrays a and b:

[[1 0 0]
 [1 1 0]
 [0 1 1]]
[[ 4.27082885  4.95718491  5.03538203]
 [ 8.41593579  3.17425233  3.99337567]
 [ 3.90636291  4.96133978  3.61849744]]

Direct output of np.where for the two conditions a nonzero and b>4:

(array([0, 1, 2], dtype=int64), array([0, 0, 1], dtype=int64))

Stacked in a human-friendly way:

[[0 0]
 [1 0]
 [2 1]]

Upvotes: 2

Related Questions