Numpy select rows based on condition

Question

I'm trying to use numpy to remove rows from a two dimensional array where the first value of the row (so the element at index 0) does not match a certain condition.

I am able to do this with regular python using two loops, but I would like to do it more efficiently with numpy, e.g. with numpy.where

I have been trying various things with numpy.where and numpy.delete but I struggle with the fact that I want to select rows by using a condition that only needs to be verified by the first element, and not the second (I dont care about the value of the second element)

Here is an example where I only want to keep the rows where the first value of each row is 6.

Input:

[[0,4],
 [0,5],
 [3,5],
 [6,8],
 [9,1],
 [6,1]]

Output:

[[6,8],
 [6,1]]

You just need to use [2D indexing](https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html). `arr = arr[arr[:,0] == 6]` — roganjosh, Sep 24 '19 at 11:23

Mad Physicist · Accepted Answer · 2021-12-04T23:10:32.270

28

Use a boolean mask:

mask = (z[:, 0] == 6)
z[mask, :]

This is much more efficient than np.where because you can use the boolean mask directly, without having the overhead of converting it to an array of indices first.

One liner:

z[z[:, 0] == 6, :]

edited Dec 04 '21 at 23:10

answered Sep 24 '19 at 11:34

Mad Physicist

107,652
25
181
264

4

We can make it even more simple one liner, as without ":" it means just the row: z=z[z[:,0]==6] – BootMaker May 26 '20 at 22:22

ravishankar chavare · Answer 2 · 2019-09-24T11:34:02.530

7

Program:

import numpy as np
np_array = np.array([[0,4],[0,5],[3,5],[6,8],[9,1],[6,1]])
rows=np.where(np_array[:,0]==6)
print(np_array[rows])

Output:

[[6 8]
 [6 1]]

And If You Want to Get Into 2d List use

np_array[rows].tolist()

Output of 2d List

[[6, 8], [6, 1]]

edited Sep 24 '19 at 11:34

answered Sep 24 '19 at 11:28

ravishankar chavare

503
3
14

Numpy select rows based on condition

2 Answers2

Linked

Related