How to generate the indices of all elements of a 3D array

Question

How should I change the below code using numpy without using foor loop?

# for loop
results = []
for i in range(256):
  for j in range(256):
    for k in range(256):
      results.append([i, j, k])

print(results)
[[  0,   0,   0],
 [  0,   0,   1],
 [  0,   0,   2],
 ...,
 [255, 255, 253],
 [255, 255, 254],
 [255, 255, 255]])

Asked By: Qimin Chen

||

Source

Answer 1

Try this:

n = 256
m = 3

result = np.vstack([
    np.tile(np.repeat(np.arange(n), n**(m-e-1)), n**e)
    for e in range(m)
]).T

It works with any dimension of your indices. Example with n=3, m=2:

array([[0, 0],
       [0, 1],
       [0, 2],
       [1, 0],
       [1, 1],
       [1, 2],
       [2, 0],
       [2, 1],
       [2, 2]])

Example with n=2, m=3:

array([[0, 0, 0],
       [0, 0, 1],
       [0, 1, 0],
       [0, 1, 1],
       [1, 0, 0],
       [1, 0, 1],
       [1, 1, 0],
       [1, 1, 1]])

Answered By: Riccardo Bucco

Answer 2

To get the indices of a 256x256x256 (3D) array (without creating it) use Numpy.indices:

dims = (256, 256, 256)
grid = np.indices(dims)
grid[0]        # row indices
grid[1]        # column indices
grid[2]        # z-axis indices

You can reshape these into a 2d array if you want:

np.indices(dims).reshape(len(dims), -1).T

array([[  0,   0,   0],
       [  0,   0,   1],
       [  0,   0,   2],
       ...,
       [255, 255, 253],
       [255, 255, 254],
       [255, 255, 255]])

Or maybe this is what you are looking for:

my_3d_array = np.random.random(dims)
for x in np.nditer(my_3d_array):
    print(x)

0.9550954204220937
0.2571918730243552
0.5651493396119187
0.4984195476733584
0.7378580503708124
... etc

If you also want the indices you can use the numpy equivalent of enumerate:

for ind, x in np.ndenumerate(my_3d_array):
    print(ind)

(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 0, 3)
(0, 0, 4)
... etc.

Answered By: Bill

Answer 3

Just discovered another way to do it.

dims = (256, 256, 256)
idx = np.stack(
    np.unravel_index(np.arange(np.product(dims)), dims)
).T

array([[  0,   0,   0],
       [  0,   0,   1],
       [  0,   0,   2],
       ...,
       [255, 255, 253],
       [255, 255, 254],
       [255, 255, 255]])

Answered By: Bill

How to generate the indices of all elements of a 3D array

Question:

Answers: