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]]

Answer 1

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, :]

Answer 2

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]]