Python - K Maximum Elements with Index in List

The list is a popular mutable data type in Python that can hold heterogeneous data. We can access the list elements by using a process called indexing. The indexing means the distance of any element from the first element; hence it starts from 0. This article will explore how to find the k maximum elements and index number. We will achieve this through several methods like brute force, recursion, heapq module, Numpy library, enumerate functions, etc. We will also explore the time and space complexity of the algorithms used.

Using The Brute Force Method

Brute force is the simplest way to solve any problem. This involves thinking about the problem with little optimization. Although this gives fair results, this needs to be better optimized regarding space and time complexity.

Example

In the following code, we have defined a function named k_max_elements which takes the list and the value k as the parameters. We iterated k times, and for each iteration, we used the "max" method of Python to find the maximum element. We appended the value and the index to an initialized list. In each iteration, we replaced the max element with '-inf' to ensure this is not counted for the next iteration.

def k_max_elements(lst, k):
    result = []
    for _ in range(k):
        max_value = max(lst)
        max_index = lst.index(max_value)
        result.append((max_value, max_index))
        lst[max_index] = float('-inf') 
    return result
lst = [10, 5, 8, 12, 3]
k = 3
result = k_max_elements(lst, k)
print(result)

Output

[(12, 3), (10, 0), (8, 2)]

Using Recursion

Recursion is a classical problem−solving technique that breaks a large problem into smaller chunks of problems. We aim to solve that smaller chunk of problems and ultimately solve the bigger problem. In a recursion solution, we also define a base case, and our program terminates when we reach the base case.

Example

In the following code, we have used the recursion method to find the k maximum elements. Our base case is k=0. We followed the top−down approach, where we tried to recursively find the maximum elements by calling the function to the base case. We appended the element and the index to an initialized list and returned it.

def k_max_elements(lst, k):
    if k == 0:
        return []
    max_value = max(lst)
    max_index = lst.index(max_value)
    lst[max_index] = float('-inf') 
    remaining_max_elements = k_max_elements(lst, k - 1)
    return [(max_value, max_index)] + remaining_max_elements
lst = [7,5,4,8,9,6,2,1]
k = 4
print(f'The original list is: {lst}')
result = k_max_elements(lst, k)
print(f"The {k} maximum elements and the index are: {result}")

Output

The original list is: [7, 5, 4, 8, 9, 6, 2, 1]
The 4 maximum elements and the index are: [(9, 4), (8, 3), (7, 0), (6, 5)]

Using The Heapq module

The 'heapq' module in Python provides us with the implementation to apply the heap algorithms. It offers functionalities for creating and manipulating heap data structures, which we can use to find the smallest or largest elements in a collection efficiently. A heap is a binary tree−like data structure in which each parent node has a value less than or equal to its child nodes. It is based on the LIFO principle − “Last In, First out”.

Example

In the following code, we used the 'heapq' library to find the k maximum elements of the list. Under the custom function, we have iterated over the list and pushed the value and the index to a heap till we have iterated k times. The heap maintains the k smallest elements at the root, so we can effectively track the k largest elements encountered so far. For elements beyond the first k, we pushed the current element−value−index tuple into the resulting heap and simultaneously popped out the smallest element from the heap. This ensures that the resulting heap always contains the k largest elements encountered.

import heapq

def k_max_elements(lst, k):
    result = []
    for i, value in enumerate(lst):
        if i < k:
            heapq.heappush(result, (value, i))
        else:
            heapq.heappushpop(result, (value, i))
    return sorted(result, reverse=True)

lst = [10, 5, 8, 12, 3]
k = 3
result = k_max_elements(lst, k)
print(f'The original list is: {lst}')
print(f"The {k} maximum elements and the index are: {result}")

Output

The original list is: [10, 5, 8, 12, 3]
The 3 maximum elements and the index are: [(12, 3), (10, 0), (8, 2)]

Using Numpy Array

Numpy is another open sources library built to deal with numeric and scientific computations. The library provides a function called “argpartition” which returns the indices of the k largest elements in the list. The function partitions the list based on the Kth largest element and returns the indices of the partitioned elements. If we are interested only in k elements, we can use the indexing property to trim out only the k elements.

Example

In the following code, we have used the “argpartition” method of Pandas to divide the list into two parts, with one part containing all the indices of k's largest elements. Next, we used the indexing property to trim only the k elements. We used the list comprehension method of Python to append the index and the corresponding elements of the indices to a list called the result. Next, we used the "sorted" method to sort the list and set the parameter 'reverse=True' to ensure it's sorted from higher to lower values.

import numpy as np

def k_max_elements(lst, k):
    indices = np.argpartition(lst, -k)[-k:]
    result = [(lst[i], i) for i in indices]
    return sorted(result, reverse=True)

lst = [4567,7,443,456,56,5467,74]
k = 4
result = k_max_elements(lst, k)
print(f'The original list is: {lst}')
print(f"The {k} maximum elements and the index are: {result}")

Output

The original list is: [4567, 7, 443, 456, 56, 5467, 74]
The 4 maximum elements and the index are: [(5467, 5), (4567, 0), (456, 3), (443, 2)]

• Time Complexity: O(n + k * log(k))

• Space complexity: O(k)

Using The Lambda And Enumerate Method

Lambda functions are special kinds of functions in Python. They are functions without any names. Lambda functions are especially useful when we want to apply a function to an iterable object and are sure we won't reuse the code elsewhere. The key advantage of using the lambda function is that we can make compact logic in one line to apply some operation. This makes the code less readable but consumes fewer lines of code.

Example

In the following example, we have used the combination of ‘sorted’, ‘enumerate’, and ‘lambda’ functions to return the k maximum elements. We passed the list and the value 'k' to the function k_max_elements. Under the function, we have used the lambda function and enumerate to associate the elements with the index. We used the sorted method to sort the list. Next, we used list comprehension to append the value and the index to the list.

def k_max_elements(lst, k):
    sorted_lst = sorted(enumerate(lst), key=lambda x: x[1], reverse=True)
    result = [(value, index) for index, value in sorted_lst[:k]]
    return result


lst = [45,7,4,46,56,5467,74]
k = 4
result = k_max_elements(lst, k)
print(f'The original list is: {lst}')
print(f"The {k} maximum elements and the index are: {result}")

Output

The original list is: [45, 7, 4, 46, 56, 5467, 74]
The 4 maximum elements and the index are: [(5467, 5), (74, 6), (56, 4), (46, 3)]

• Time Complexity: O(n * log(n))

• Space complexity: O(n)

Using The Pandas Libray

Pandas is a well−known open−sourced library of Python. The library is designed to deal with data preprocessing and manipulations. Pandas mainly deals with Pandas series and Pandas data frames. It provides the sort_values method to sort all the values from highest to lowest. We can utilize the method to find the k maximum elements. Since Pandas deals with the data frame, we can deal with multiple columns at a time, and hence we can easily keep track of both the index and the elements.

Example

In the following example, we first imported the Pandas library. Next, we created the function k_max_elements, which takes the list and the value 'k' as the parameters and returns the k maximum elements. We used the Dataframe method of Pandas to create a data frame containing the index and the value of the list's elements. Next, we used the 'nlargest' and the 'sort_+values' methods to find the 'k' largest elements in sorted form. We used the zip method to unpack the value and indices of the elements and returned the result.

import pandas as pd

def k_max_elements(lst, k):
    df = pd.DataFrame({'value': lst, 'index': range(len(lst))})
    df = df.nlargest(k, 'value').sort_values('index')
    result = list(zip(df['value'], df['index']))
    return result


lst = [45,7,4,46,56,5467,74]
k = 4
result = k_max_elements(lst, k)
print(f'The original list is: {lst}')
print(f"The {k} maximum elements and the index are: {result}")

Output

The original list is: [45, 7, 4, 46, 56, 5467, 74]
The 4 maximum elements and the index are: [(46, 3), (56, 4), (5467, 5), (74, 6)]

• Time Complexity: O(n * log(k))

• Space complexity: O(n)

Conclusion

This article taught us how to find k maximum elements with the index in a list. Python offers several ways to deal with this problem. We can use the recursion, loop statements, etc., to build our custom logic. However, Python also provides us with many in−built methods which can assist us in performing the same. Numpy, Pandas, and Heapq are examples of libraries and packages which enable us to perform this. However, the readers should note that these are not the only methods. We can perform several tweakings, including changing the output to other data structures.