How to select specific rows and columns in 2d array

Question

If I have an 2d array such as

A = np.arange(16).reshape(4,4)

How can I select row = [0, 2] and column = [0, 2] using parameters? In MATLAB, I can simply do A[row, column] but in python this will select 2 elements corresponding to (0,0) and (2,2).

Is there anyway I can do this using some parameters as in MATLAB? The output should be like [0 2

8 10]

Answer 1

You can use the following

A = np.arange(16).reshape(4,4)
print np.ravel(A[row,:][:,column])

to get:

array([ 0,  2,  8, 10])

Answer 2

MATLAB creates a 2D mesh when indexed with vectors across dimensions. So, in MATLAB, you would have -

A =
     0     1     2     3
     4     5     6     7
     8     9    10    11
    12    13    14    15
>> row = [1, 3]; column = [1, 3];
>> A(row,column)
ans =
     0     2
     8    10

---

Now, in NumPy/Python, indexing with the vectors across dimensions selects the elements after making tuplets from each element in those vectors. To replicate the MATLAB behaviour, you need to create a mesh of such indices from the vectors. For the same, you can use np.meshgrid -

In [18]: A
Out[18]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

In [19]: row = [0, 2]; column = [0, 2];

In [20]: C,R = np.meshgrid(row,column)

In [21]: A[R,C]
Out[21]: 
array([[ 0,  2],
       [ 8, 10]])

Answer 3

To select a block of elements - as MATLAB does, the 1st index has to be column vector. There are several ways of doing this:

In [19]: A = np.arange(16).reshape(4,4)
In [20]: row=[0,2];column=[0,2]

In [21]: A[np.ix_(row,column)]
Out[21]: 
array([[ 0,  2],
       [ 8, 10]])
In [22]: np.ix_(row,column)
Out[22]: 
(array([[0],
        [2]]), array([[0, 2]]))

In [23]: A[[[0],[2]],[0,2]]
Out[23]: 
array([[ 0,  2],
       [ 8, 10]])

The other answer uses meshgrid . We could probably list a half dozen variations.

Good documentation in this section: http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#purely-integer-array-indexing