pythonarraysnumpyn-dimensional

comparing values in numpy array of n-dimension


Need to compare each value inside a numpy array and return 1 to the largest value and 0 to the others. I am having problem with different numbers of [].

Input Example:

[[[0.6673975 0.33333233]]
.
.
.
[[0.33260247 0.6673975]]]

Expected Output:

[[[1 0]]
.
.
.
[[0 1]]]

Solution

  • Max over axis:

    If, as suggested by Joe in the comments, you're looking for a maximum along an axis, then, for axis axis,

    np.moveaxis((np.moveaxis(ar, axis, 0) == ar.max(axis)).astype(int), 0, axis)
    

    or, a bit faster,

    (ar == np.broadcast_to(np.expand_dims(ar.max(axis), axis), ar.shape)).astype(int)
    

    Should cover the n-dimensional case.

    Ex:

    ar = np.random.randint(0, 100, (2, 3, 4))
    
    ar
    Out[157]: 
    array([[[17, 28, 22, 31],
            [99, 51, 65, 65],
            [46, 24, 93,  4]],
    
           [[ 5, 84, 85, 79],
            [ 7, 80, 27, 25],
            [46, 80, 90,  3]]])
    
    (ar == np.broadcast_to(np.expand_dims(ar.max(-1), -1), ar.shape)).astype(int)
    Out[159]: 
    array([[[0, 0, 0, 1],
            [1, 0, 0, 0],
            [0, 0, 1, 0]],
    
           [[0, 0, 1, 0],
            [0, 1, 0, 0],
            [0, 0, 1, 0]]])
    ar.max(-1)
    Out[160]: 
    array([[31, 99, 93],
           [85, 80, 90]])
    

    Max over full array:

    On the off-chance you're trying to identify elements equal to the maximum over the whole array,

    (ar == ar.max()).astype(int)
    

    should give what you're looking for.