Итерация по всем одномерным подмассивам многомерного массива

Question

Какой самый быстрый способ перебора всех одномерных подмассивов n-мерного массива в python.

Например, рассмотрим трехмерный массив:

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

желаемая последовательность выходов от итератора:

a[:,0,0]
a[:,0,1]
..
a[:,2,3]
a[0,:,0]
..
a[1,:,3]
a[0,0,:]
..
a[1,2,:]

Answer 1

Вот компактная реализация такого итератора:

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

Это приведет к появлению подмассивов в порядке, указанном вами в вашем сообщении:

for x in iter1d(a):
    print x

print

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

Хитрость заключается в том, чтобы перебирать все оси и для каждой оси преобразовывать массив в двумерный массив, строки которого являются искомыми одномерными подмассивами.

Answer 2

Может быть более эффективный способ, но это должно сработать ...

import itertools
import numpy as np

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

colon = slice(None)
dimensions = [range(dim) + [colon] for dim in a.shape]

for dim in itertools.product(*dimensions):
    if dim.count(colon) == 1:
        print a[dim]

Это дает (я опускаю тривиальный бит кода для печати левой части этого ...):

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

Ключевым моментом здесь является то, что индексирование a с (например) a[0,0,:] эквивалентно индексации с a[(0,0,slice(None))]. (Это просто обычная нарезка на Python, ничего специфичного для numpy. Чтобы доказать это самим, вы можете написать фиктивный класс всего с __getitem__ и распечатать то, что передается, когда вы индексируете экземпляр вашего фиктивного класса.).

Итак, нам нужны все возможные комбинации от 0 до nx, от 0 до ny, от 0 до nz и т. Д. И None для каждой оси.

Однако нам нужны одномерные массивы, поэтому нам нужно отфильтровать все, что больше или меньше, чем один None (т.е. нам не нужны a[:,:,:], a[0,:,:], a[0,0,0] и т. Д.).

Надеюсь, в любом случае это имеет какой-то смысл ...

Редактировать: Я предполагаю, что точный порядок не имеет значения ... Если вам нужен точный порядок, который вы указали в своем вопросе, вам нужно изменить это ...

Answer 3

Ваши друзья - это объекты slice(), метод numpy's ndarray.__getitem__() и, возможно, itertools.chain.from_iterable.