如何使用python將所有索引值的總和加到numpy ndarray的每個元素中？

Question

我有一個numpy ndarray，讓我們舉個例子（但可以完全不同）：

[[[0 0 0]
  [1 1 1]
  [0 0 0]]

 [[1 0 1]
  [1 0 1]
  [1 0 1]]]

 [[1 0 0]
  [1 1 0]
  [1 1 1]]]

我想將每個維度的索引總和添加到每個元素。 結果將是：

[[[ 0  1  2]
  [ 4  5  6]
  [ 6  7  8]]

 [[10 10 12]
  [13 13 15]
  [16 16 18]]

 [[19 19 20]
  [22 23 23]
  [25 26 27]]]

為此，我構建了另一個ndarray：

shp = a.shape
b = np.arange(shp[0]**len(shp)).reshape(shp)

我得到了我的結果：

result = a+b

我想知道是否有一個更直接的解決方案，它不需要創建第二個ndarray，一種在“位置”上執行相同操作的方法？

Answer 1

只需創建開放式網格即可，基本上是一維陣列，擴展到更多的暗點，然后利用broadcasting將其添加到輸入陣列中-

m,n,r = a.shape
I,J,K = np.ogrid[:m,:n,:r]
out = a + I*n*r + J*r + K

因此，就內存占用而言，與range-based相比，我們只創建了9（= m + n + r）個元素，而不是27（= m * n * r）個元素。

樣品運行-

In [41]: a
Out[41]: 
array([[[0, 0, 0],
        [1, 1, 1],
        [0, 0, 0]],

       [[1, 0, 1],
        [1, 0, 1],
        [1, 0, 1]],

       [[1, 0, 0],
        [1, 1, 0],
        [1, 1, 1]]])

In [42]: m,n,r = a.shape

In [43]: I,J,K = np.ogrid[:m,:n,:r]

In [44]: a + I*n*r + J*r + K
Out[44]: 
array([[[ 0,  1,  2],
        [ 4,  5,  6],
        [ 6,  7,  8]],

       [[10, 10, 12],
        [13, 13, 15],
        [16, 16, 18]],

       [[19, 19, 20],
        [22, 23, 23],
        [25, 26, 27]]])

N維陣列盒

對於通用n維數組，將a添加到自身中-

shp = a.shape
grid = np.ogrid[tuple(map(slice, shp))]
scale = np.r_[np.array(shp)[::-1].cumprod()[::-1][1:],1]
for i,j in zip(grid,scale):
    a += i*j