0

I have a 3D numpy array A of shape 10 x 5 x 3. I also have a vector B of length 3 (length of last axis of A). I want to compare each A[:,:,i] against B[i] where i = 0:2 and replace all values A[:,:,i] > B[i] with B[i].

Is there a way to achieve this without a for loop.

Edit: I tried the argmax across i = 0:2 using a for loop python replace values in 2d numpy array

Community
  • 1
  • 1
Zanam
  • 4,607
  • 13
  • 67
  • 143

1 Answers1

3

You can use numpy.minimum to accomplish this. It returns the element-wise minimum between two arrays. If the arrays are different sizes (such as in your case), then the arrays are automatically broadcast to the correct size prior to comparison.

A = numpy.random.rand(1,2,3)
# array([[[ 0.79188   ,  0.32707664,  0.18386629],
#         [ 0.4139146 ,  0.07259663,  0.47604274]]])

B = numpy.array([0.1, 0.2, 0.3])

C = numpy.minimum(A, B)
# array([[[ 0.1       ,  0.2       ,  0.18386629],
#         [ 0.1       ,  0.07259663,  0.3       ]]])

Or as suggested by @Divakar if you want to do in-place replacement:

numpy.minimum(A, B, out=A)
Community
  • 1
  • 1
Suever
  • 64,497
  • 14
  • 82
  • 101