How to create Numpy array using fancy indexing

Question

How do use numpy's fancy indexing to create this, I would like the fastest performance:

array([[ 1,  2,  3,  4, 16, 31],
       [ 2,  3,  4,  5, 17, 32],
       [ 3,  4,  5,  6, 18, 33],
       [ 4,  5,  6,  7, 19, 34],
       [ 5,  6,  7,  8, 20, 35],
       [ 6,  7,  8,  9, 21, 36],
       [ 7,  8,  9, 10, 22, 37],
       [ 8,  9, 10, 11, 23, 38],
       [ 9, 10, 11, 12, 24, 39],
       [10, 11, 12, 13, 25, 40]]

Starting with this:

   a   = np.arange(0,10)
   aa  = np.arange(0,50)
   y   = 1
   AA  = [(aa[np.array([x+y, 1+x+y, 2+x+y, 3+x+y,  15+x+y, 30+x+y])]) for x, i in enumerate(a)]

I get this,

[array([ 2,  3,  4,  5, 17, 32]),
 array([ 3,  4,  5,  6, 18, 33]),
 array([ 4,  5,  6,  7, 19, 34]),
 array([ 5,  6,  7,  8, 20, 35]),
 array([ 6,  7,  8,  9, 21, 36]),
 array([ 7,  8,  9, 10, 22, 37]),
 array([ 8,  9, 10, 11, 23, 38]),
 array([ 9, 10, 11, 12, 24, 39]),
 array([10, 11, 12, 13, 25, 40]),
 array([11, 12, 13, 14, 26, 41])]   

Answer 1

Making use of the given variables a, aa and y, here's one using broadcasting for the outer addition -

offset = np.array([0,1,2,3,15,30])
out = aa[a[:,None] + offset + y]

Using add-ufunc's explicit outer method -

out = aa[np.add.outer(a , offset + y)]

---

Out of bounds case

For out of bounds case (aa being smaller than required), we can pad zeros with aa and then index into it -

offset = np.array([0,1,2,3,15,30])
idx = np.add.outer(a , offset + y)
aa_p = np.pad(aa,(0,idx.max()-len(a)+1), 'constant')
out = aa_p[idx]

Or initialize output array and then create a mask of valid places for assigning -

offset = np.array([0,1,2,3,15,30])
idx = np.add.outer(a , offset + y)
mask = idx < len(aa)
out = np.zeros(idx.shape, dtype=aa.dtype)
out[mask] = aa[idx[mask]]

Sample input, output -

In [234]: a   = np.arange(0,10)
     ...: aa  = np.arange(4,42)
     ...: y   = 6
     ...: 

In [235]: out
Out[235]: 
array([[10, 11, 12, 13, 25, 40],
       [11, 12, 13, 14, 26, 41],
       [12, 13, 14, 15, 27,  0],
       [13, 14, 15, 16, 28,  0],
       [14, 15, 16, 17, 29,  0],
       [15, 16, 17, 18, 30,  0],
       [16, 17, 18, 19, 31,  0],
       [17, 18, 19, 20, 32,  0],
       [18, 19, 20, 21, 33,  0],
       [19, 20, 21, 22, 34,  0]])

Answer 2

We can make use of broadcasting here:

>>> np.arange(0,10).reshape(-1,1) + np.array([*range(1,5),16,31])
array([[ 1,  2,  3,  4, 16, 31],
       [ 2,  3,  4,  5, 17, 32],
       [ 3,  4,  5,  6, 18, 33],
       [ 4,  5,  6,  7, 19, 34],
       [ 5,  6,  7,  8, 20, 35],
       [ 6,  7,  8,  9, 21, 36],
       [ 7,  8,  9, 10, 22, 37],
       [ 8,  9, 10, 11, 23, 38],
       [ 9, 10, 11, 12, 24, 39],
       [10, 11, 12, 13, 25, 40]])

Here we create a 10×1 matrix that ranges from 0 to (excluding) 10, and we create a 1× 6 matrix with data [1,2,3,4,16,31].

In case you want y to be the "offset", you can write it as:

>>> y = 1
>>> np.arange(y,y+10).reshape(-1,1) + np.array([*range(0,4),15,30])
array([[ 1,  2,  3,  4, 16, 31],
       [ 2,  3,  4,  5, 17, 32],
       [ 3,  4,  5,  6, 18, 33],
       [ 4,  5,  6,  7, 19, 34],
       [ 5,  6,  7,  8, 20, 35],
       [ 6,  7,  8,  9, 21, 36],
       [ 7,  8,  9, 10, 22, 37],
       [ 8,  9, 10, 11, 23, 38],
       [ 9, 10, 11, 12, 24, 39],
       [10, 11, 12, 13, 25, 40]])