Это упражнение в вещании:
>>> import numpy as np
>>> from itertools import islice
>>>
>>> A = np.arange(27).reshape(3,3,3)
>>> def per(n):
... while True:
... yield np.random.permutation(n)
...
>>> pg = per(3)
>>>
>>> p0 = next(pg)[..., None, None]
>>> p1 = np.array([p for p in islice(pg, 3)])[..., None]
>>> p2 = np.array([[p for p in islice(pg, 3)] for _ in range(3)])
>>>
>>> p0
array([[[2]],
[[1]],
[[0]]])
>>> p1
array([[[1],
[0],
[2]],
[[1],
[0],
[2]],
[[0],
[2],
[1]]])
>>> p2
array([[[1, 0, 2],
[0, 2, 1],
[0, 1, 2]],
[[2, 1, 0],
[1, 2, 0],
[2, 1, 0]],
[[1, 2, 0],
[2, 1, 0],
[2, 1, 0]]])
>>> A[p0, p1, p2]
array([[[22, 21, 23],
[18, 20, 19],
[24, 25, 26]],
[[14, 13, 12],
[10, 11, 9],
[17, 16, 15]],
[[ 1, 2, 0],
[ 8, 7, 6],
[ 5, 4, 3]]])
Общее время:
import numpy as np
from itertools import islice
def per(n=None):
while True:
n = (yield n if n is None else np.random.permutation(n)) or n
def sweep(shp):
pg = per()
pg.send(None)
redshp = len(shp)*[1]
sz = 1
for j, k in enumerate(shp):
pg.send(k)
redshp[j] = k
yield np.reshape((*islice(pg, sz),), redshp)
sz *= k
# example
a = np.arange(24).reshape((2,3,4))
idx = *sweep(a.shape),
for i in idx:
print(i)
print(a[idx])