NumPy array indexing

Question

I want to extract the second and the 3rd to the fifth columns of the NumPy array, how would I go about it?

A = array([[0, 1, 2, 3, 4, 5, 6], [4, 5, 6, 7, 4, 5, 6]])
A[:, [1, 4:6]]

This obviously doesn't work.

Answer 1

Assuming I've understood you -- it's usually a good idea to explicitly specify the output you want, because it's not obvious -- you could use numpy.r_:

In [27]: A
Out[27]: 
array([[0, 1, 2, 3, 4, 5, 6],
 [4, 5, 6, 7, 4, 5, 6]])
In [28]: A[:, [1,3,4,5]]
Out[28]: 
array([[1, 3, 4, 5],
 [5, 7, 4, 5]])
In [29]: A[:, r_[1, 3:6]]
Out[29]: 
array([[1, 3, 4, 5],
 [5, 7, 4, 5]])
In [37]: A[1:, r_[1, 3:6]]
Out[37]: array([[5, 7, 4, 5]])

which you can then flatten or reshape as you like. r_ is basically a convenience function to generate the right indices, e.g.

In [30]: r_[1, 3:6]
Out[30]: array([1, 3, 4, 5])

• tchakravarty: Quick question, why is the index array not inclusive?

Answer 2

Perhaps you are looking for this?

In [10]: A[1:, [1]+range(3,6)]
Out[10]: array([[5, 7, 4, 5]])

Note this gives you the second, fourth, fifth and six columns of all rows but the first.

Answer 3

The second element is A[:,1]. Elements 3-5 (I'm assuming you want inclusive) are A[:,2:5]. You won't be able to extract them with a single call. To get them as an array, you could do

import numpy as np
A = np.array([[0, 1, 2, 3, 4, 5, 6], [4, 5, 6, 7, 4, 5, 6]])
my_cols = np.hstack((A[:,1][...,np.newaxis], A[:,2:5]))

The np.newaxis stuff is just to make A[:,1] a 2D array, consistent with A[:,2:5].

Hope this helps.