python - 在二维数组中插入元素的有效方法

Question

我未能成功将此函数转换为矢量化函数:

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)

输出:

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

有什么帮助吗？

Answer 1

您可以切换到阵列的一维 View a:

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

重新计算一维数组的索引:

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

插入一维数组

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

并将结果 reshape 回二维

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

所以，最后我们得到了

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

@askewchan 在评论中提出的在线用户，使用 np.ravel_multi_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]])