Selecting specific rows and columns from NumPy array

Question

I ve been going crazy trying to figure out what stupid thing I m doing wrong here   I m using NumPy  and I have specific row indices and specific column indices that I want to select from  Here s the gist of my problem   import numpy as np  a   np arange 20  reshape  5 4     array    0   1   2   3              4   5   6   7              8   9  10  11             12  13  14  15             16  17  18  19       If I select certain rows  it works print a  0  1  3        array    0   1   2   3              4   5   6   7             12  13  14  15       If I select certain rows and a single column  it works print a  0  1  3   2    array   2   6  14      But if I select certain rows AND certain columns  it fails print a  0 1 3    0 2     Traceback  most recent call last       File   lt stdin gt    line 1  in  lt module gt    ValueError  shape mismatch  objects cannot be broadcast to a single shape   Why is this happening  Surely I should be able to select the 1st  2nd  and 4th rows  and 1st and 3rd columns  The result I m expecting is   a  0 1 3    0 2     gt    0   2                          4   6                          12  14

User · Accepted Answer

Fancy indexing requires you to provide all indices for each dimension. You are providing 3 indices for the first one, and only 2 for the second one, hence the error. You want to do something like this:

>>> a[[[0, 0], [1, 1], [3, 3]], [[0,2], [0,2], [0, 2]]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

That is of course a pain to write, so you can let broadcasting help you:

>>> a[[[0], [1], [3]], [0, 2]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

This is much simpler to do if you index with arrays, not lists:

>>> row_idx = np.array([0, 1, 3])
>>> col_idx = np.array([0, 2])
>>> a[row_idx[:, None], col_idx]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

User · Answer

USE     gt  gt  gt  a  0 1 3      0 2   array    0   2        4   6       12  14      OR    gt  gt  gt  a  0 1 3    2  array    0   2        4   6       12  14

User · Answer

Using np ix  is the most convenient way to do it  as answered by others   but here is another interesting way to do it    gt  gt  gt  rows    0  1  3   gt  gt  gt  cols    0  2    gt  gt  gt  a rows  T cols  T  array    0   2            4   6           12  14

User · Answer

As Toan suggests  a simple hack would be to just select the rows first  and then select the columns over that    gt  gt  gt  a  0 1 3                   Returns the rows you want array    0   1   2   3            4   5   6   7           12  13  14  15     gt  gt  gt  a  0 1 3          0 2      Selects the columns you want as well array    0   2            4   6           12  14       Edit  The built-in method  np ix   I recently discovered that numpy gives you an in-built one-liner to doing exactly what  Jaime suggested  but without having to use broadcasting syntax  which suffers from lack of readability   From the docs      Using ix  one can quickly construct index arrays that will index the   cross product  a np ix   1 3   2 5    returns the array   a 1 2  a 1 5     a 3 2  a 3 5       So you use it like this    gt  gt  gt  a   np arange 20  reshape  5 4    gt  gt  gt  a np ix   0 1 3    0 2    array    0   2            4   6           12  14      And the way it works is that it takes care of aligning arrays the way Jaime suggested  so that broadcasting happens properly    gt  gt  gt  np ix   0 1 3    0 2    array   0            1            3     array   0  2       Also  as MikeC says in a comment  np ix  has the advantage of returning a view  which my first  pre-edit  answer did not  This means you can now assign to the indexed array    gt  gt  gt  a np ix   0 1 3    0 2      -1  gt  gt  gt  a     array   -1   1  -1   3           -1   5  -1   7            8   9  10  11           -1  13  -1  15           16  17  18  19

[python] Selecting specific rows and columns from NumPy array

Examples related to python

Examples related to arrays

Examples related to numpy

Examples related to multidimensional-array

Examples related to numpy-slicing