Home > OS >  numpy apply "where" on one column?
numpy apply "where" on one column?

Time:11-06

[[11.  5.]
 [24.  6.]
 [39. 12.]
 [14.  1.]
 [25. 12.]]

That's my matrix, now I want to apply the following condition to it: 10 < a < 15 but only to the first column

As result I want an array with the indices [0. 3.]

My attempts with "where" failed, numpy applied it to both columns

CodePudding user response:

The first column is a[:, 0]. np.flatnonzero is a version of np.nonzero (which is where with only one argument) that returns a flat array instead of a tuple of indices.

You can do something like

mask = (10 > a[:, 0]) & (a[:, 0] < 15)
idx = np.flatnonzero(mask)

Another way is

idx, = np.nonzero(mask)

Or even

idx, = np.where(mask)

Notice the comma after idx in the last two examples, which triggers argument unpacking.

  • Related