CODEWITHSUNDEEP
pythonarraysnumpy
Cupitor
Cupitor

Reputation: 11637

Efficient way of inserting elements in a 2D array

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?

Upvotes: 1

Views: 192

Answers (2)

askewchan
askewchan

Reputation: 46530

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

Upvotes: 1

alko
alko

Reputation: 48317

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]])

Upvotes: 2

Related Questions