Extracting multiple sets of rows/ columns from a 2D numpy array-CodePudding

I have a 2D numpy array from which I want to extract multiple sets of rows/ columns.

# img is 2D array
img = np.arange(25).reshape(5,5)
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

I know the syntax to extract one set of row/ column. The following will extract the first 4 rows and the 3rd and 4th column as shown below

img[0:4, 2:4]
array([[ 2,  3],
       [ 7,  8],
       [12, 13],
       [17, 18]])

However, what is the syntax if I want to extract multiple sets of rows and/or columns? I tried the following but it leads to an invalid syntax error

img[[0,2:4],2]

The output that I am looking for from the above command is

array([[ 2],
       [12],
       [17]])

I tried searching for this but it only leads to results for one set of rows/ columns or extracting discrete rows/ columns which I know how to do. Any help on this will be highly appreciated.

CodePudding user response：

IIUC, you can use numpy.r_ to generate the indices from the slice:

img[np.r_[0,2:4][:,None],2]

output:

array([[ 2],
       [12],
       [17]])

intermediates:

np.r_[0,2:4]
# array([0, 2, 3])

np.r_[0,2:4][:,None]  # variant: np.c_[np.r_[0,2:4]]
# array([[0],
#        [2],
#        [3]])

CodePudding user response：

Easiest way to do this is create your indexes as a list, then slice with the list:

i_list = [0,2,3]
img[i_list,2]

output:

array([ 2, 12, 17])