Replacing values in a NumPy array with is_max

Question

I have an array like this:

a = np.array(
    [[
        [
            [0.        , 0.        ],
            [0.        , 0.        ],
            [0.        , 0.        ]
        ],
        [
            [0.07939843, 0.13330124],
            [0.20078699, 0.17482429],
            [0.49948007, 0.52375553]
        ]
    ]]
)

a.shape
>>> (1, 2, 3, 2)

Now I need to compare the values in the last dimension and replace the maximum with 1 and the minimum with 0. So I need a result like this:

array([[[[1., 0.],
          [1., 0.],
          [1., 0.]],

         [[0., 1.],
          [1., 0.],
          [0., 1.]]]])

Well, I have tried a nested loop like this:

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        for k in range(a.shape[2]):
            if a[i, j, k, 0] > a[i, j, k, 1]:
                a[i, j, k, 0] = 1
                a[i, j, k, 1] = 0
            else:
                a[i, j, k, 0] = 0
                a[i, j, k, 1] = 1

However, I need a faster way, maybe a built-in function of NumPy library.

@Sujay I checked that. Thank you. But that seems to be useful when we want to compare the values with a constant value. How can I compare two values in the same dimension using the approach you suggested? — A. Arash Chitgar
– A. Arash Chitgar, Commented Aug 22, 2021 at 17:43

akuiper · Accepted Answer · 2021-08-22 17:45:20Z

1

Manually construct the indices for 1 using argmax:

i, j, k, _ = a.shape
idx = np.ogrid[:i,:j,:k] + [a.argmax(-1)]
a[:] = 0
a[tuple(idx)] = 1

a
array([[[[1., 0.],
         [1., 0.],
         [1., 0.]],

        [[0., 1.],
         [1., 0.],
         [0., 1.]]]])

answered Aug 22, 2021 at 17:45

akuiper

216k33 gold badges363 silver badges380 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Replacing values in a NumPy array with is_max

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related