How to select 3d array elements with another 3-d array of coordinates?

Question

Let's say I have two array, one with all zero elements and the second one with indices that should be filled with 1. This can be done with the following code:

A = np.zeros((100, 50))
B = np.concatenate([np.random.randint(low=0, high=99, size=(10, 1)),
                    np.random.randint(low=0, high=49, size=(10, 1))],axis=1)
A[B[:, 0], B[:, 1]] = 1

However, this get trickier with adding another dimension. Now my A array is of shape (6, 100, 50) and my B array of coordinates is of shape (6, 10, 2):

A = np.zeros((6, 100, 50))

B = []
for i in range(6):
    B0 = np.concatenate([np.random.randint(low=0, high=99, size=(10, 1)),
                    np.random.randint(low=0, high=49, size=(10, 1))],axis=1)
    B.append(B0)
B = np.stack(B)

How can I then select elements of A with coordinates stored in B? First dimension is the same and contain the coordinates for respective matrices stored in first dimension of A.

Answer 1

Indexing specific values in a 3D array has to be done with 3 values, so the positions of subarrays in B is frankly kind of useless. There are 2 ways that I'm aware of to approach getting the index of the subarray position.

One way is to use a for loop to iterate through each subarray in B, but personally, I prefer explicitly creating a column in B with the first dimension index (though I recognize that this could eat up memory if B is very large). Here's how adding a column and indexing with that would work:

B_shape = B.shape
num_dimensions = 3
first_dimension = B_shape[0]
second_dimension = B_shape[1]
third_dimension = B_shape[2]

# Create a new column in B stating the index of the first dimension,
# rather than relying on the position as a stand-in. Flatten B so there are no subarrays,
# only rows of 3.
indices = np.repeat(np.arange(first_dimension)[:,None,None], second_dimension, axis=1)
B2 = np.concatenate((indices, B), axis=2).reshape(first_dimension*second_dimension, third_dimension + 1)

# Index A using the new B2 array.
A2 = A[tuple(B2.T)]

Let me know if you have any questions.