Функция должна возвращать 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])
Но это работает не так как я хочу.
