# can I split numpy array with mask?

You can generate the b as a broadcasted sum between the indices you want, and a shift-vector. Then you can broadcast again into a bigger size. Since the output in your examples are not depending on the a array, I'm disregarding that.

from numpy import array, broadcast_to, arange
from numpy.random import random

a = random((10,10,10)) # not used on the code at all.... don't understand what it is for...

b = [0,2,3]
b_array = array(b)
b_shifts = arange(3).reshape(-1,1)
c_cell= b+b_shifts # here, they are broadcasted toegether. one is a row-vector and one is a column-vector...



you might want to create b_shifts with some other method depending on step size and so on...

from numpy import array, arange
a = arange(2*2*10).reshape((2,2,10)) # some example input
b = array([0,2,3])                   # the 'template' to extract
shifts = arange(3).reshape(-1,1)     # 3 is the number of repeats
indexer = b+shifts                   # broadcasted sum makes a matrix
c = a[:,:,indexer]                   # extract



This will take the b array as a kind of template, and repeat it with a certain shift. Finally, it will extract those entries from every array a[i,j,:] into c[i,j,:,:]. Otput from above is:

print(a)

[[[ 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 25 26 27 28 29]
[30 31 32 33 34 35 36 37 38 39]]]


print(c)

[[[[ 0  2  3]
[ 1  3  4]
[ 2  4  5]]
[[10 12 13]
[11 13 14]
[12 14 15]]]
[[[20 22 23]
[21 23 24]
[22 24 25]]
[[30 32 33]
[31 33 34]
[32 34 35]]]]


A moving windows approach using as_strided:

In [1]: a = np.arange(6)
In [2]: a
Out[2]: array([0, 1, 2, 3, 4, 5])
In [3]: as_strided = np.lib.stride_tricks.as_strided


For this one-shift the strides parameter is easy. The shape requires more thinking - how many rows we expect, and the maximum index:

In [5]: b = as_strided(a, shape=(3,4), strides=(8,8))
In [6]: b
Out[6]:
array([[0, 1, 2, 3],
[1, 2, 3, 4],
[2, 3, 4, 5]])


Then select columns:

In [8]: b[:,[0,2,3]]
Out[8]:
array([[0, 2, 3],
[1, 3, 4],
[2, 4, 5]])


To extend it to the 3d case, I'll work from https://stackoverflow.com/a/60881930/901925, LudvigH's answer

In [10]: a = np.arange(2*2*10).reshape((2,2,10)) # some example input
...: b = np.array([0,2,3])
In [11]: a
Out[11]:
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, 25, 26, 27, 28, 29],
[30, 31, 32, 33, 34, 35, 36, 37, 38, 39]]])
In [12]: a.shape
Out[12]: (2, 2, 10)
In [13]: a.strides
Out[13]: (160, 80, 8)
In [18]: a1 = as_strided(a, shape=(2,2,3,4), strides=(160,80,8,8))
In [19]: a1
Out[19]:
array([[[[ 0,  1,  2,  3],
[ 1,  2,  3,  4],
[ 2,  3,  4,  5]],

[[10, 11, 12, 13],
[11, 12, 13, 14],
[12, 13, 14, 15]]],

[[[20, 21, 22, 23],
[21, 22, 23, 24],
[22, 23, 24, 25]],

[[30, 31, 32, 33],
[31, 32, 33, 34],
[32, 33, 34, 35]]]])


This is just an extension of the first case, with the first 2 dimensions just going-along for the ride. It's the last dimensions that expanded into the 2d window.

Again selecting a subset of the columns:

In [20]: a1[:,:,:,b]
Out[20]:
array([[[[ 0,  2,  3],
[ 1,  3,  4],
[ 2,  4,  5]],

[[10, 12, 13],
[11, 13, 14],
[12, 14, 15]]],

[[[20, 22, 23],
[21, 23, 24],
[22, 24, 25]],

[[30, 32, 33],
[31, 33, 34],
[32, 34, 35]]]])


The as_strided step is efficient, creating a view. But the indexing will make a copy. We'd have to do some timing to test it against LudvigH's shifted index approach.