Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Iteration through all 1 dimensional subarrays of a multi-dimensional array

What is the fastest way to iterate through all one dimensional sub-arrays of an n dimensional array in python.

For example consider the 3-D array:

import numpy as np 
a = np.arange(24)
a = a.reshape(2,3,4)

The desired sequence of yields from the iterator is :

a[:,0,0]
a[:,0,1]
..
a[:,2,3]
a[0,:,0]
..
a[1,:,3]
a[0,0,:]
..
a[1,2,:]
like image 216
fodon Avatar asked Apr 16 '11 15:04

fodon


1 Answers

Here is a compact implementation of such an iterator:

def iter1d(a):
    return itertools.chain.from_iterable(
        numpy.rollaxis(a, axis, a.ndim).reshape(-1, dim)
        for axis, dim in enumerate(a.shape))

This will yield the subarrays in the order you gave in your post:

for x in iter1d(a):
    print x

prints

[ 0 12]
[ 1 13]
[ 2 14]
[ 3 15]
[ 4 16]
[ 5 17]
[ 6 18]
[ 7 19]
[ 8 20]
[ 9 21]
[10 22]
[11 23]
[0 4 8]
[1 5 9]
[ 2  6 10]
[ 3  7 11]
[12 16 20]
[13 17 21]
[14 18 22]
[15 19 23]
[0 1 2 3]
[4 5 6 7]
[ 8  9 10 11]
[12 13 14 15]
[16 17 18 19]
[20 21 22 23]

The trick here is to iterate over all axes, and for each axis reshape the array to a two-dimensional array the rows of which are the desired one-dimensional subarrays.

like image 187
Sven Marnach Avatar answered Oct 06 '22 23:10

Sven Marnach