Как красиво сгенерировать 4-х мерный массив, содержащий в 4-м измерении индексы по первым 3-м измерениям?

Question

[nirvimel]

Функция должна возвращать 3+1 мерный массив, в четвертом измерении в эллементах 0, 1, 2 должны лежать значения равные индексам, соответственно 1-го, 2-го и 3-го измерения так, чтобы выполнялось равенство data[x, y, z] == array([x, y, z]).
Я написал так:

def array_of_indexes(x, y, z, dtype=int):
    array = np.empty((x, y, z, 3), dtype=dtype)
    array[:, :, :, 0] = np.reshape(np.arange(x), (x, 1, 1))
    array[:, :, :, 1] = np.reshape(np.arange(y), (y, 1))
    array[:, :, :, 2] = np.arange(z)
    return array

Это работает, но мне не нравится создание 5 временных массивов (каждый reshape создает еще один) и три прохода по результирующму массиву.
Я думаю должно существовать решение проще и красивее. Например, на основе meshgrid. numpy.meshgrid и numpy.dstack намекают на такое решение:

def array_of_indexes_alt(x, y, z):
    return np.dstack(np.mgrid[0:x, 0:y, 0:z])

Но это работает не так как я хочу.

Answer 1

[Alexander Movchan]

def array_of_insexes(x, y, z):
    return [
               [
                   [
                       (i, j, k) for k in range(z)
                   ] for j in range(y)
               ] for i in range(x)
           ]

Comments

[nirvimel]

Речь о векторных вычислениях (numpy).
Проход в циклах по всему массиву - не тот порядок скоростей.