Python – K Maximum elements with Index in List

GIven a List, extract K Maximum elements with their indices.

Input : test_list = [5, 3, 1, 4, 7, 8, 2], K = 2 Output : [(4, 7), (5, 8)] Explanation : 8 is maximum on index 5, 7 on 4th. Input : test_list = [5, 3, 1, 4, 7, 10, 2], K = 1 Output : [(5, 10)] Explanation : 10 is maximum on index 5.

Method #1 : Using sorted() + index()

The combination of above functions provide a way of finding solution to this problem. In this, we initially perform sort and extract K maximum elements, then are encapsulated in tuple with their ordering in original list.

Python3

# Python3 code to demonstrate working of
# K Maximum elements with Index in List
# Using sorted() + index()

# initializing list
test_list = [5, 3, 1, 4, 7, 8, 2]

# printing original list
print("The original list : " + str(test_list))

# initializing K
K = 3

# Using sorted() + index()
# using sorted() to sort and slice K maximum elements
temp = sorted(test_list)[-K:]
res = []
for ele in temp:

    # encapsulating elements with index using index()
    res.append((test_list.index(ele), ele))

# printing result
print("K Maximum with indices : " + str(res))

Output

The original list : [5, 3, 1, 4, 7, 8, 2]
K Maximum with indices : [(0, 5), (4, 7), (5, 8)]

Time Complexity: O(nlogn), where n is the elements of list
Auxiliary Space: O(n), where n is the size of list

Method #2 : Using enumerate() + itemgetter()

The combination of above functions can be used to solve this problem. In this, we perform the task of getting indices using enumerate() and itemgetter() is used to get the elements. 

Python3

# Python3 code to demonstrate working of
# K Maximum elements with Index in List
# Using enumerate() + itemgetter()
from operator import itemgetter

# initializing list
test_list = [5, 3, 1, 4, 7, 8, 2]

# printing original list
print("The original list : " + str(test_list))

# initializing K
K = 3

# Using enumerate() + itemgetter()
# Making index values pairs at 1st stage
res = list(sorted(enumerate(test_list), key = itemgetter(1)))[-K:]

# printing result
print("K Maximum with indices : " + str(res))

Output

The original list : [5, 3, 1, 4, 7, 8, 2]
K Maximum with indices : [(0, 5), (4, 7), (5, 8)]

Time Complexity: O(n) where n is the number of elements in the list “test_list”. 
Auxiliary Space: O(n) where n is the number of elements in the list “test_list”. 

Method 3 : using the heapq module in Python.

step-by-step approach

1. Import the heapq module.
2. Initialize a heap with the first K elements of the list.
3. Traverse the remaining elements of the list and for each element, compare it with the smallest element in the heap. If the element is greater than the smallest element, replace the smallest element with the current element.
4. Once all elements have been processed, the heap will contain the K maximum elements.
5. For each element in the heap, find its index in the original list using the index() method.
6. Combine the index and the element as a tuple, and add it to the result list.
7. Return the result list.

Python3

import heapq

# initializing list
test_list = [5, 3, 1, 4, 7, 8, 2]

# printing original list
print("The original list : " + str(test_list))

# initializing K
K = 3

# using heapq module
heap = test_list[:K]
heapq.heapify(heap)

for i in range(K, len(test_list)):
    if test_list[i] > heap[0]:
        heapq.heappop(heap)
        heapq.heappush(heap, test_list[i])

# finding indices of K maximum elements
res = []
for elem in heap:
    index = test_list.index(elem)
    res.append((index, elem))

# printing result
print("K Maximum with indices : " + str(res))

Output

The original list : [5, 3, 1, 4, 7, 8, 2]
K Maximum with indices : [(0, 5), (4, 7), (5, 8)]

 The time complexity of this approach is O(n log K), where n is the length of the list and K is the number of maximum elements to be found.

The auxiliary space used by this approach is O(K), as we are only storing K elements in the heap at any given time.