How to replace "array elements" of a numpy.ndarray with respect to a condition without usi-CodePudding

I have a three dimensional numpy.ndarray with shape (5, 4, 3) and I would like to replace some of the arrays with three elements inside my array by another one. For instance, I would like to replace all [3, 3, 3] elements by [0, 0, 0]. Below, I present a solution I found but this solution is not efficient enough and can be really long on big arrays.

I have created the following numpy.ndarray:

c = np.array([[[3,3,3],[2,2,2],[3,3,3],[4,4,4]],[[1,1,1],[2,2,2],[7,3,3],[4,4,4]],[[1,1,1],[3,3,3],[3,3,3],[4,4,4]],[[1,1,1],[2,2,2],[3,8,3],[3,3,3]],[[3,3,3],[2,2,2],[3,3,3],[4,4,4]]])

This gives me:

>>> c

array([[[3, 3, 3],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [7, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [3, 3, 3],
        [3, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [3, 8, 3],
        [3, 3, 3]],

       [[3, 3, 3],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]]])

>>> c.shape

    (5, 4, 3)

I would like to replace all the elements [3, 3, 3] by [0, 0, 0].

The solution I found is the following:

# I reshape my array to get only a 2D array
c_copy = c.reshape(c.shape[0] * c.shape[1], c.shape[2])

# Then I loop through all the elements of the array to replace [3, 3, 3] by [0, 0, 0]
c_copy[[np.array_equal(e, [3, 3, 3]) for e in c_copy]] = [0,0,0]

# And I reshape my copy to get the original shape back
c_modified = c_copy.reshape(c.shape)

This works well:

>>> c_modified

array([[[0, 0, 0],
        [2, 2, 2],
        [0, 0, 0],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [7, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [0, 0, 0],
        [0, 0, 0],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [3, 8, 3],
        [0, 0, 0]],

       [[0, 0, 0],
        [2, 2, 2],
        [0, 0, 0],
        [4, 4, 4]]])

However the loop for e in c_copy is killing me. This is fine for this small ndarray but I have an array of 9000000 elements and I haven't found any efficient solution.

What could I do to speed the computation?

Live example

CodePudding user response：

You need find index of [3,3,3] for this you can use all(axis=-1), then replace with [0,0,0]:

row, col = np.where((c==3).all(-1))
c[row, col] = [0,0,0]
print(c)

Output:

[[[0 0 0]
  [2 2 2]
  [0 0 0]
  [4 4 4]]

 [[1 1 1]
  [2 2 2]
  [7 3 3]
  [4 4 4]]

 [[1 1 1]
  [0 0 0]
  [0 0 0]
  [4 4 4]]

 [[1 1 1]
  [2 2 2]
  [3 8 3]
  [0 0 0]]

 [[0 0 0]
  [2 2 2]
  [0 0 0]
  [4 4 4]]]