Return Array Index Based on Highest Value of Another Index-CodePudding

I have a multidimensional array where I want to return the 1st index, of the highest value of the 2nd index if they have the same zeroth indexes. For example:

 array = [[ 0, 2, 54]
          [ 0, 3, 83]
          [ 0, 4, 92]
          [ 1, 5, 52]
          [ 1, 6, 28]
          [ 2, 7, 20]
          [ 2, 8, 22]

I want it to return: 4,5,8 (Basically, the zeroth indexes are the categories the row belongs to, the 1st index is what I want it to return, based on the highest value in the 2nd indexes of the same zeroth category)

I'm not sure how to even start so apologies for not including what I have tried.

CodePudding user response：

As an answer has been posted, we set a few things up:

from itertools import groupby
from operator import itemgetter

array = [[ 0, 2, 54],
         [ 0, 3, 83],
         [ 0, 4, 92],
         [ 1, 5, 52],
         [ 1, 6, 28],
         [ 2, 7, 20],
         [ 2, 8, 22]]

We can then do this with a list comprehension one-liner:

>>> [max(v, key=third)[1] for _, v in groupby(sorted(array, key=itemgetter(0)), itemgetter(0))]
[4, 5, 8]
>>>

First we sort the two-dimensional list by the first element in each sublist. Then we group them by the first element. Now we iterate over those groups, finding the max sublist in each group based on the third element, and we extract the second element from each.

CodePudding user response：

You can use a hashmap (dictionary) to keep track of the maxes and the second indices to do this in O(n) time. Try this:

array = [[ 0, 2, 54],
          [ 0, 3, 83],
          [ 0, 4, 92],
          [ 1, 5, 52],
          [ 1, 6, 28],
          [ 2, 7, 20],
          [ 2, 8, 22]]

def get_max_second_indices(array):
    # We are going to populate this in the format {zero_index: {second_index: index, max_third: max}}      
    final_values = {}
    for sublist in array:
        # Check if the final_values contains the zero index
        if final_values.get(sublist[0], None) is None:
            final_values[sublist[0]] = {"second_index": sublist[1], "max": sublist[2]}
        else:
            if sublist[2] > final_values[sublist[0]]['max']:
                final_values[sublist[0]] = {"second_index": sublist[1], "max": sublist[2]}
    result = []
    for i in sorted(final_values):
        result.append(final_values[i]['second_index'])
    return result
    
print(get_max_second_indices(array))

(returns [4,5,8])