在numpy中插入3維矩陣中的非對齊元素

Question

我正在使用3維矩陣使用numpy 1.9和python 2.7.5。 下面是一個例子：

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

我想從B A的所有元素插入，根據C. 實施例：c[0] = 1 => After A[0, 1, :] has to be inserted 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，它與我的真的很相似，只是解決方案只適用於二維矩陣，而我正在使用三維。

這裏是我的解決辦法，但我得到不正確的結果：

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

編輯：原來的答案包括全面瓦解/重塑 回來，但在3D，你需要很多照顧做是正確的。有一個更容易的解決方案，考慮到np.insert和 np.take接受「軸」參數並允許插入多值的事實。 這仍然需要一些重塑，但它不訴諸於 np.choose。另外，還要注意的mi+1參數np.insert到後插入 ，而不是選擇行之前：

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

Answer 2

這似乎爲您（非工作）矢量解決方案的最後一行替換工作：

linear = np.insert(A.ravel(), index_map + r2[::-1], B.ravel()[index_map - 1]) 
linear.reshape(A.shape[0], A.shape[1] + 1, A.shape[2]) 

這就好比你的矢量化解決方案，但有一對夫婦的調整，以獲得正確的索引。第一個關鍵是要認識到我需要「撤銷」你的index_map的1。下一個頓悟是當你插入linear時，你需要抵消每一行中的索引，因爲當你插入元素時，後面的元素會被移回。因此，雖然index_map是[4,5,6,22,23,24]，但我們實際上需要[6,6,6,24,24,24]，並且我剛剛爲此目的而重用/濫用了r2[::-1]。

似乎B.ravel()[index_map - 1]可以簡化爲B[r1,C2,r2]。和多一點的簡化，消除r2[::-1]怪異減爲您提供：

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