Changing values of a 3D array based on 2D coordinates (Python)

Question

I have 2 arrays. Call them 'A' and 'B'. The shape of array 'A' is (10,10,3), and the shape of array 'B' is (10,10).

Now, I have acquired the coordinates of certain elements of array 'B' as a list of tuples. Let's call the list 'TupleList'. I want to now make the values of all elements in array 'A' equal to 0, except for those elements present at the coordinates in 'TupleList'.

Imagine the arrays are image arrays. Where A being an RGB image has the 3rd dimension.

How can I do so?

I am having trouble doing this because of the extra 3rd dimension that array 'A' has. Otherwise, it is quite straightforward with the use of np.where(), given that I know the limited number of values that array 'B' can take.

score 1 · Accepted Answer · answered Aug 18 '22 at 12:28

Here is a solution by transposing the indices list into list of indices for the first dimension and list of indices for the second dimension:

import numpy as np

nrows, ncols = 5, 5

arr = np.arange(nrows * ncols * 3, dtype=float).reshape(nrows, ncols, 3)
idxs = [(0, 1), (1, 3), (3, 2), (4, 2)]

idxs_dim0, idxs_dim1 = zip(*idxs)

res = np.zeros(arr.shape)
res[idxs_dim0, idxs_dim1] = arr[idxs_dim0, idxs_dim1]

Checking the result:

from itertools import product

for idx_dim0, idx_dim1 in product(range(nrows), range(ncols)):
    value = res[idx_dim0, idx_dim1]
    if (idx_dim0, idx_dim1) in idxs:
        assert np.all(value == arr[idx_dim0, idx_dim1])
    else:
        assert np.all(value == (0, 0, 0))

Changing values of a 3D array based on 2D coordinates (Python)

1 Answers1