CODEWITHSUNDEEP
pythonarrays
Chelsea Zou
Chelsea Zou

Reputation: 27

Return Array Index Based on Highest Value of Another Index

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.

Upvotes: 0

Views: 87

Answers (3)

Chris
Chris

Reputation: 36660

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=itemgetter(2)[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.

Upvotes: 2

AboAmmar
AboAmmar

Reputation: 5559

A simple method using a dictionary. Store the 1st and 2nd indices in the zeroth index key of the dictionary. Compare the 2nd indices and replace, if necessary, for each zeroth index.

d = {}
for i, j, k in array:
    if i in d: 
        if d[i][1] < k:
            d[i] = j, k
    else:
        d[i] = j, k

This will store the maxima in the j values of the dictionary.

[j for j, k in d.values()] 
 # [4, 5, 8]

Upvotes: 0

M. Chak
M. Chak

Reputation: 532

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])

Upvotes: 1

Related Questions