按条件提取numpy数组中的特定列

Question

I have a homework assignment to extract a 2-dimensional numpy array out of another 2-dimensional np array by choosing specific columns by condition (not by range).

So I have an array A with shape (3, 50000) . I am trying to get a new array with shape (3, x) for some x < 50000 with the original columns of A that satisfy the third cell in the column is -0.4 < z < 0.1`.

For example if:

A = [[1,2,3],[2,0.5,0],[9,-2,-0.2],[0,0,0.5]]

I wish to have back:

B = [[2,0.5,0],[9,-2,-0.2]

I have tried to make a bool 1 rank array that holds true on the columns I want, and to some how combine between the two. The problem it's output is 1 rank array which is not what I am looking for. And I got some ValueErrors..

bool_idx = (-0.4 < x_y_z[2] < 0.1)

This code made some troubles:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I can do it with some loops but NumPy got so many beautiful function I am sure I am missing something here..

Answer 1

In Python, the expression -0.4 < x_y_z[2] < 0.1 is roughly equivalent to -0.4 < x_y_z[2] and x_y_z[2] < 0.1 . The and operator decides the truth value of each part of the expression by converting it into a bool. Unlike Python lists and tuples, numpy arrays do not support the conversion.

The correct way to specify the condition is with bitwise & (which is unambiguous and non-short-circuiting), rather than the implicit and (which short circuits and is ambiguous in this case):

condition = ((x_y_z[2, :] > - 0.4) & (x_y_z[2, :] < 0.1))

condition is a boolean mask that selects the columns you want. You can select the rows with a simple slice:

selection = x_y_z[:, condition]