Insert array into all places in another array

Question

For two arrays, say, a = np.array([1,2,3,4]) and b = np.array([5,6]), is there a way, if any, to obtain a 2d array of the following form without looping:

[[5 6 1 2 3 4]
 [1 5 6 2 3 4]
 [1 2 5 6 3 4]
 [1 2 3 5 6 4]
 [1 2 3 4 5 6]]

i.e. to insert b in all possible places of a.

And if loops are unavoidable, how to do it the most computationally efficient way (a can be long, the length of b is irrelevant)?

Example of how it can be done using loops is trivial:

a = np.array([1,2,3,4])
b = np.array([5,6])
rows = len(a) + 1
cols = len(a) + len(b)
res = np.empty([rows, cols])
for i in range(rows):
    res[i, i:len(b)+i] = b
    res[i, len(b)+i:] = a[i:]
    res[i, 0:i] = a[0:i]
print(rows.astype(int))

[[5 6 1 2 3 4]
 [1 5 6 2 3 4]
 [1 2 5 6 3 4]
 [1 2 3 5 6 4]
 [1 2 3 4 5 6]]

Answer 1

Consider using numba acceleration. This happens to be what numba is best at. For your example, it can speed up nearly 6 times:

from timeit import timeit

import numpy as np
from numba import njit


a = np.arange(1, 5)
b = np.array([5, 6])


def fill(a, b):
    rows = a.size + 1
    cols = a.size + b.size
    res = np.empty((rows, cols))
    for i in range(rows):
        res[i, i:b.size + i] = b
        res[i, b.size + i:] = a[i:]
        res[i, :i] = a[:i]
    return res


if __name__ == '__main__':
    print('before:', timeit(lambda: fill(a, b)))
    fill = njit(fill)
    print('after:', timeit(lambda: fill(a, b)))

Output:

before: 9.488150399993174
after: 1.6149254000047222