How can I add to each element of a numpy ndarray the sum of all its index values with python ?

Question

I have a numpy ndarray, let's take an example (but it can be totally different):

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

I want to add to each element the sum of its indexes in all dimensions. The result here would be:

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

To do so, I built another ndarray:

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

And I got my result:

result = a+b

I would like to know if there is a more direct solution, which wouldn't need the creation of this second ndarray, a way to do the same operation 'on location' ??

Answer 1

Simply create open grids, that are basically 1D arrays extended to more dims and add into input array leveraging broadcasting -

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

Hence, in terms of memory occupancy, we are creating only 9 (=m+n+r) more elements as opposed to 27 (= m * n * r) elements with the range-based soltuion.

Sample run -

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-dim array case

For a generic n-dim array, a to add into itself -

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