CODEWITHSUNDEEP
pythonnumpy
justyy
justyy

Reputation: 6041

How to Extract Columns from Numpy

I have a numpy array which is 300 rows and 5 columns 

X[X[:,0]==1,[1,2]]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-154-de5e74bc1a0b> in <module>()
----> 1 X[X[:,0]==1,[1,2]]

ValueError: shape mismatch: objects cannot be broadcast to a single shape

How to make this work? I want to filter the first column (equals 1) but returns a copy of second and third column.

Upvotes: 0

Views: 318

Answers (1)

Alex Riley
Alex Riley

Reputation: 177048

It's simplest to separate the two indexing cases. Look at column one first, select the appropriate rows, and then pick out columns 1 and 2 from that array:

>>> a = np.random.randint(0, 2, (3, 5))
>>> a
array([[0, 0, 0, 0, 1],
       [1, 0, 1, 0, 1],
       [0, 0, 1, 1, 0]])

>>> a[a[:,0] == 1][:,[1,2]]
array([[0, 1]])

The code in the question, a[a[:,0] == 1, [1,2]], look syntactically similar but is doing something different.

For example a[a[:,0] == 0, [1,2]] is, in the case of a, equivalent to a[[0,2], [1,2]]. This indexing picks out exactly two elements from a: the element at row 0, column 1 and the element at row 2, column 2. It does not select the rows and columns of the array.

Upvotes: 1

Related Questions