python - 在 numpy 的 3 维矩阵中插入非对齐元素

Question

我正在使用 numpy 1.9 和 python 2.7.5 处理 3 维矩阵。这是一个例子:

>>> A
array([[[ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.]],

       [[ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.]]])

>>> B
array([[[-1., -1., -1.],
        [99., 100., 101.],
        [-1., -1., -1.],
        [-1., -1., -1.],
        [-1., -1., -1.]],

       [[-1., -1., -1.],
        [-1., -1., -1.],
        [102., 103., 104.],
        [-1., -1., -1.],
        [-1., -1., -1.]]])

>>> C
array([1, 2])

根据 C，我想在 A 中插入来自 B 的所有元素。示例:c[0] = 1 => 在 A[0, 1, :] 之后必须插入 B[0, 1, :]

这是一个预期结果的例子

>>> D
array([[[1.,  1.,  1.],
        [1.,  1.,  1.],
        [99., 100., 101.],
        [1.,  1.,  1.],
        [1.,  1.,  1.],
        [1.,  1.,  1.]],

       [[1.,  1.,  1.],
        [1.,  1.,  1.],
        [1.,  1.,  1.],
        [102., 103., 104.]
        [1.,  1.,  1.],
        [1.,  1.,  1.]]])

我找到了 this stackoverflow question这与我的非常相似，只是解决方案仅适用于 2 维矩阵，而我正在使用 3 维矩阵。

这是我的解决方案，但我得到的结果不正确:

C2 = np.repeat(C, 3)
r1 = np.repeat(np.arange(A.shape[0]), 3)
r2 = np.tile(np.arange(3), A.shape[0])
index_map = np.ravel_multi_index((r1, C2, r2), A.shape) + 1
np.insert(A.ravel(), index_map, B.ravel()[index_map]).reshape(A.shape[0], A.shape[1] + 1, A.shape[2])

这是一个使用 for 循环的正确但缓慢的解决方案:

A_2 = np.zeros((A.shape[0], A.shape[1] + 1, A.shape[2]))
for j in xrange(np.size(C, 0)):
  i = C[j]
  A_2[j, :, :] = np.concatenate((A[j, 0:i + 1, :], [B[j, i, :]], A[j, i + 1:, :]))

有什么想法吗？

谢谢!

Answer 1

你的代码的问题是，当你需要插入几个元素顺序，你需要将它们插入到相同的位置。比较:

In [139]: x = np.ones(5); x
Out[139]: array([ 1.,  1.,  1.,  1.,  1.])

In [140]: np.insert(x, [1,2,3], 100)
Out[140]: array([   1.,  100.,    1.,  100.,    1.,  100.,    1.,    1.])

In [141]: np.insert(x, [1,1,1], 100)
Out[141]: array([   1.,  100.,  100.,  100.,    1.,    1.,    1.,    1.])

编辑: 原始答案包括完整的拆解/ reshape 返回，但在 3d 中，您需要非常小心才能正确完成。有一个考虑到 np.insert 和np.take 接受“轴”参数并允许插入多值。这仍然需要一些 reshape ，但它不会求助于np.选择。另外，请注意 np.insert 的 mi+1 参数以插入在之后，而不是在所选行之前:

In [50]: mi = np.ravel_multi_index([np.arange(A.shape[0]), C], A.shape[:2]); mi
Out[50]: array([1, 7])

In [51]: bvals = np.take(B.reshape(-1, B.shape[-1]), mi, axis=0); bvals
Out[51]: 
array([[  99.,  100.,  101.],
       [ 102.,  103.,  104.]])

In [52]: result = (np.insert(A.reshape(-1, A.shape[2]), mi + 1, bvals, axis=0)
                   .reshape(A.shape[0], -1, A.shape[2])); result
Out[52]: 
array([[[   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [  99.,  100.,  101.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.]],

       [[   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [ 102.,  103.,  104.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.]]])

这是原来的答案:

In [18]: ixs = np.repeat(np.array([np.arange(A.shape[0]),
                                    C+1,
                                    np.zeros(A.shape[0], dtype=np.int_)]),
                          A.shape[2], axis=1); ixs
   ....: 
Out[18]: 
array([[0, 0, 0, 1, 1, 1],
       [2, 2, 2, 3, 3, 3],
       [0, 0, 0, 0, 0, 0]])

In [19]: mi = np.ravel_multi_index(ixs, A.shape); mi
Out[19]: array([ 6,  6,  6, 24, 24, 24])

In [20]: result = (np.insert(A.ravel(), mi, bvals)
                    .reshape(A.shape[0], A.shape[1] +1, A.shape[2])); result
   ....: 
Out[20]: 
array([[[   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [  99.,  100.,  101.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.]],

       [[   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [ 102.,  103.,  104.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.]]])

In [21]: result = (np.insert(A.ravel(), mi, bvals)
                    .reshape(A.shape[0], A.shape[1] +1, A.shape[2])); result
   ....: 
Out[21]: 
array([[[   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [  99.,  100.,  101.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.]],

       [[   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.],
        [ 102.,  103.,  104.],
        [   1.,    1.,    1.],
        [   1.,    1.,    1.]]])