How to Extract Multiple Columns from NumPy 2D Matrix?


The numpy package is a powerful toolkit for Python. With numpy you can easily do matrix (like Matlab), plot and do other data/numbers/statistics. Suppose you have a two dimensional array (also treated as matrix, i.e. the array of vectors):

arr = np.random.randint(0, 5, (10, 5))
array([[2, 0, 1, 0, 3],
       [1, 3, 4, 0, 1],
       [2, 0, 1, 1, 0],
       [2, 4, 0, 1, 2],
       [1, 1, 3, 0, 0],
       [0, 1, 4, 3, 1],
       [0, 3, 4, 4, 1],
       [1, 3, 0, 2, 1],
       [4, 1, 1, 4, 4],
       [3, 0, 0, 2, 0]])

And, it is easy if you want to extract the second and third column and return a new copy:

arr[:,[1,2]]
array([[0, 1],
       [3, 4],
       [0, 1],
       [4, 0],
       [1, 3],
       [1, 4],
       [3, 4],
       [3, 0],
       [1, 1],
       [0, 0]])

However, if you want to filter the rows at the same time, e.g., only the first column equals to 2.

arr[arr[:,0]==2,[1,2]]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
ipython-input-16-4f6b52878ed7 in module()
---- 1 arr[arr[:,0]==2,[1,2]]

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

Although, this works (return all columns)

arr[arr[:,0]==2,:]
array([[2, 0, 1, 0, 3],
       [2, 0, 1, 1, 0],
       [2, 4, 0, 1, 2]])

The correct way is to first select the rows and then return the wanted columns:

arr[arr[:,0]==2,:][:,[1,2]]
array([[0, 1],
       [0, 1],
       [4, 0]])

Two deep-copies will be made.

–EOF (The Ultimate Computing & Technology Blog) —

GD Star Rating
a WordPress rating system
216 words
Last Post: How to Compare Version Number String in C#?
Next Post: How to Plot Scatter in Numpy Python?

The Permanent URL is: How to Extract Multiple Columns from NumPy 2D Matrix?

Leave a Reply