Home > Blockchain >  Numpy return sub-array with minimal sum
Numpy return sub-array with minimal sum

Time:04-03

Given a 2D numpy array, e.g. np.array([[1,1,0], [1,2,0], [0,0,4]]), I want to return the sub-array with the minimal sum.

In this case, this would be the array np.array([1,1,0]).

How can I achieve this, ideally only using built-in numpy functions? It seems numpy does not allow specifying a key function, unlike vanilla python.

CodePudding user response:

First we can use the build in function numpy.sum and get sums allong axis 1. How a numpy array is indexed can be found here.

    array = np.array([[1,1,0], [1,2,0], [0,0,4]])
    internal_sums = np.sum(array, axis=1)

Now we have all the sums stored in the variable 'internal_sums'. Next we use the numpy function 'numpy.argmin' to get the index of the minimal sum.

    minsum_index = np.argmin(internal_sums)

Now we simply get the array at the correct index

    minsum_array = array[minsum_index]
  • Related