Search code examples
pythonlistmaxsublist

Finding out indices of maxima of sublists in python

I've a list x of n sublists. I want to make a new_list of length n containing the indices of maximum of each sublist. How can I pull this off?

For example:

x = [[1,4,2,3,5],[3,5,2,1,3],[5,2,1,4,3],[1,2,3,5,4],[4,1,5,2,3]]

(I've taken the maxima the same in each sublist for the ease of use)

And the output must hence be:

new_list = [4,1,0,3,2]
Solution

  • If you have access to numpy it is very easy:

    In [9]: import numpy as np

In [10]: x = [[1,4,2,3,5],[3,5,2,1,3],[5,2,1,4,3],[1,2,3,5,4],[4,1,5,2,3]]

In [11]: a = np.array(x) # Assumes all sublists are of the same length

In [12]: np.argmax(a, axis=1)
Out[12]: array([4, 1, 0, 3, 2])