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.

Check https://stackoverflow.com/questions/19666626/replace-all-elements-of-python-numpy-array-that-are-greater-than-some-value — , Aug 22 '21 at 17:32
@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, Aug 22 '21 at 17:43

score 1 · Accepted Answer · answered Aug 22 '21 at 17:45

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.]]]])

Replacing values in a NumPy array with is_max

1 Answers1