Replace elements in a numpy ndarray accessed via index array

Question

I want to replace certain columns in a multi-dimensional array with the value that I have. I tried the following.

cols_to_replace = np.array([1, 2, 2, 2])

original_array = np.array([[[255, 101,  51],
    [255, 101, 153],
    [255, 101, 255]],

   [[255, 153,  51],
    [255, 153, 153],
    [255, 153, 255]],

   [[255, 203,  51],
    [255, 204, 153],
    [255, 205, 255]],

   [[255, 255,  51],
    [255, 255, 153],
    [255, 255, 255]]], dtype=int)

Replace just the cols with (0, 0, 255)

I was hoping that I can index all the columns using the array cols_to_replace

 original_array[:, cols_to_replace] = (0, 0, 255)

This gave a wrong answer!

Unexpected output.

array([[[255, 101,  51],
    [  0,   0, 255],
    [  0,   0, 255]],

   [[255, 153,  51],
    [  0,   0, 255],
    [  0,   0, 255]],

   [[255, 203,  51],
    [  0,   0, 255],
    [  0,   0, 255]],

   [[255, 255,  51],
    [  0,   0, 255],
    [  0,   0, 255]]])

My expected output is

array([[[255, 101,  51],
    [  0,   0, 255],
    [255, 101, 255]],

   [[255, 153,  51],
    [255, 153, 153],
    [  0,   0, 255]],

   [[255, 203,  51],
    [255, 204, 153],
    [  0,   0, 255]],

   [[255, 255,  51],
    [255, 255, 153],
    [  0,   0, 255]]])

• What is really happening?

• How do I accomplish what I am trying to do (that is access, col 1, col 2, col 2, col 2 in each of these rows and replace the values.

• If I want to delete those columns, is there a numpy way to do it?

Answer 1

Your expected output is produced by:

>>> original_array[np.arange(cols_to_replace.size), cols_to_replace] = 0, 0, 255

This differs from your original approach because advanced indexing and slice indexing are evaluated "separately". By changing : to an arange we switch the zeroth dimension to advanced indexing, so that cols_to_replace is paired element-by-element with 0, 1, 2, ... in the zeroth coordinate.

You can delete your selection using a mask like so:

>>> mask = np.ones(original_array.shape[:2], bool)
>>> mask[np.arange(cols_to_replace.size), cols_to_replace] = False
>>> original_array[mask].reshape(original_array.shape[0], -1, original_array.shape[2])
array([[[255, 101,  51],
        [255, 101, 255]],

       [[255, 153,  51],
        [255, 153, 153]],

       [[255, 203,  51],
        [255, 204, 153]],

       [[255, 255,  51],
        [255, 255, 153]]])