Index two sets of columns in an array

Question

I am trying to slice columns out of an array and assign to a new variable, like so.

array1 = array[:,[0,1,2,3,15,16,17,18,19,20]]

Is there a short cut for something like this?

I tried this, but it threw an error:

array1 = array[:,[0:3,15:20]]

This is probably really simple but I can't find it anywhere.

Answer 1

Use np.r_:

Translates slice objects to concatenation along the first axis.

import numpy as np
arr = np.arange(100).reshape(5, 20)
cols = np.r_[:3, 15:20]

print(arr[:, cols])
[[ 0  1  2 15 16 17 18 19]
 [20 21 22 35 36 37 38 39]
 [40 41 42 55 56 57 58 59]
 [60 61 62 75 76 77 78 79]
 [80 81 82 95 96 97 98 99]]

At the end of the day, probably only a little less verbose than what you have now, but could come in handy for more complex cases.

Answer 2

For most simple cases like this, the best and most straightforward way is to use concatenation:

array1 = array[0:3] + array[15:20]

For more complicated cases, you'll need to use a custom slice, such as NumPy's s_, which allows for multiple slices with gaps, separated by commas. You can read about it here.

Also, if your slice follows a pattern (i.e. get 5, skip 10, get 5 etc), you can use itertools.compress, as explained by user ncoghlan in this answer.