如何使用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