David - 4 months ago 8x
Python Question

# Iterating and selecting a specific array from a multidimensional array in Python

Imagine I have something like this:

`````` import numpy as np
arra = np.arange(16).reshape(2, 2, 4)
``````

which gives

`````` array([[[0, 1, 2, 3],
[4, 5, 6, 7]],
[[8, 9, 10, 11],
[12, 13, 14, 15]]])
``````

and I want to make a loop that runs along specific subarrays (in this case, I want to run along each 'column' of each 'matrix') and sum the result (that is, selecting 0 and 4 and summing them (4), selecting 1 and 5 and summing them (6), ..., selecting 3 and 7 and summing them (10), selecting 8 and 12 and summing them (20), ..., selecting 11 and 15 and summing them (26)).

I had thought doing that with the apparently natural:

`````` for i in arra[i, j, k]:
for j in arra[i, j, k]:
for k in arra[i, j, k]:
sum...
``````

The problem is that Python certainly doesn't allow to do what I want in this way. If it were a 2D array it would be easier as I know that the iterator first runs through the rows, so you can transpose to run along the columns, but for a multidimensional (3D in this case) array (N, M, P) with N, M, P >> 1, I was wondering how it could be done.

EDIT: This question has a continuation here: Choosing and iterating specific sub-arrays in multidimensional arrays in Python

You can use `map` to get this done:

``````import numpy as np
arra = np.arange(16).reshape(2, 2, 4)
``````

Then the command:

``````map(sum, arra)
``````

gives you the desired output:

``````[array([ 4,  6,  8, 10]), array([20, 22, 24, 26])]
``````

Alternatively, you can also use a list comprehension:

``````res = [sum(ai) for ai in arra]
``````

Then `res` looks like this:

``````[array([ 4,  6,  8, 10]), array([20, 22, 24, 26])]
``````

EDIT:

If you want to add identical rows - as you mentioned in the comments below this answer - you can do (using `zip`):

``````map(sum, zip(*arra))
``````

which gives you the desired output:

``````[array([ 8, 10, 12, 14]), array([16, 18, 20, 22])]
``````

For the sake of completeness also the list comprehension:

``````[sum(ai) for ai in zip(*arra)]
``````

which gives you the same output.