pythonnumpynumpy-ndarray

NumPy Built-In Function to Find Largest Vector in an Matrix


I have a 2D numpy matrix:

arr = np.array([(1, 2), (6, 0), (3, 3), (5, 4)])

I am trying to get the output:

[5, 4]

I have tried to do the following:

max_arr = np.max(arr, axis=0)

But this finds the largest value in each column of the matrix, regardless of the vectors from which those values come from.

Maybe there is some function that will consider the sum of each vector in the array? The idea is to use only the optimized NumPy functionality in the least steps possible.


Solution

  • Compute the sum per row, then get the position of the maximum with argmax and use this for indexing:

    out = arr[arr.sum(axis=1).argmax()]
    

    Output: array([5, 4])

    Intermediates:

    arr.sum(axis=1)
    # array([3, 6, 6, 9])
    
    arr.sum(axis=1).argmax()
    # 3
    

    variant: all maxes

    If you can have multiple maxima, you can use instead:

    out = arr[(s:=arr.sum(axis=1)) == s.max()]
    

    Or for python < 3.8:

    s = arr.sum(axis=1)
    out = arr[s == s.max()]
    

    Output:

    array([[5, 4]])
    

    Note the 2D output, you would get multiple rows if several have the same maximum sum.