Efficient way of inserting elements in a 2D array

Question

I am unsuccessful in turning this function into a vectorised one:

a=np.asarray([[1,2,3],[3,4,5]])
inds=np.asarray([0,2])
vals=np.asarray([10,12])
def new_insert(arr,inds,vals):
    ret=np.zeros((arr.shape[0],arr.shape[1]+1))
    for i in range(arr.shape[0]):
        ret[i]=np.insert(arr[i],inds[i],vals[i])
    return ret
print new_insert(a,inds,vals)

With output:

[[ 10.   1.   2.   3.]
 [  3.   4.  12.   5.]]

Any helps?

Answer 1

You can switch to a 1d view of your array a :

shape = a.shape
a.shape = np.multiply(*shape)

recalculate indexes for 1-d array:

ind1d = [i+e*shape[0] for i, e in enumerate(ind)]

insert in 1d array

b = np.insert(a, ind1d, vals)

and reshape result back to 2d

b.shape = (shape[0], shape[1]+1)

So, finally, we get

>>> b
array([[10,  1,  2,  3],
       [ 3,  4, 12,  5]])

An onliner, proposed by @askewchan in comments, using np.ravel_multi_index helper function to flatten index:

>>> np.insert(a.flat, np.ravel_multi_index((np.arange(ind.size), ind), 
...               a.shape), vals).reshape(a.shape[0], -1)
array([[10,  1,  2,  3],
       [ 3,  4, 12,  5]])

Answer 2

Figured I'd post my comment to @alko's answer as an answer, since it looks a bit confusing as one line:

b = np.insert(a.flat, np.ravel_multi_index((np.arange(ind.size), ind), a.shape), vals).reshape(a.shape[0], -1)

This is basically the same as @alko's but it has a few advantages:

• It does not modify a itself, by using the a.flat iterator instead of actually changing the shape of a .
• Avoids potential bugs by using the np.ravel_multi_index to create the ind1d array instead of doing it manually.
• It is a tiny bit (10%) faster.

In steps similar to alko's, this is what it does:

ind1d = np.ravel_multi_index((np.arange(ind.size), ind), a.shape)

where ind refers to column index, so use np.arange to refer to row index. Then, insert into the a.flat iterator instead of the reshaped a :

b = np.insert(a.flat, ind1d, vals)

Finally, reshape:

b = b.reshape(a.shape[0], -1) # the -1 allows any shape at the end