Python - Find the indices for k Smallest elements

Sometimes, while working with Python lists, we can have a problem in which we wish to find the indices for K smallest elements. This task can occur in many domains such as web development and while working with Databases. We might sometimes, require to just find the indices of them. Let’s discuss a certain way to find indices of K smallest elements in list. 

Method 1: Using sorted() + lambda + list slicing 

This task can be performed using a combination of the above functions. In this the sorted(), can be used to get the container in a way that requires to get K smallest elements at front end and then the indices can be computed using list slicing. 

# Python3 code to demonstrate working of
# Smallest K elements indices
# using sorted() + lambda + list slicing

# Initialize list
test_list = [5, 6, 10, 4, 7, 1, 19]

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

# initialize K
K = 3

# Smallest K elements indices
# using sorted() + lambda + list slicing
res = sorted(range(len(test_list)), key=lambda sub: test_list[sub])[:K]

# Printing result
print("Indices list of min K elements is : " + str(res))

The original list is : [5, 6, 10, 4, 7, 1, 19]
Indices list of min K elements is : [5, 3, 0]

Time complexity: O(nlogn) (sorting using sorted() function)
Auxiliary space: O(n) (creating a sorted list of indices)

Method 2: using the heapq module   

In this method to obtain the indices of the smallest K elements in a list can be to use the heapq module in Python. The heapq module provides the implementation of heap queue algorithm. We can use the nsmallest() function from the heapq module to find the smallest K elements and their indices in a list.

import heapq

# Initialize list
test_list = [5, 6, 10, 4, 7, 1, 19]

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

# Initialize K
K = 3

# Smallest K elements indices using heapq module
smallest_indices = heapq.nsmallest(
    K, range(len(test_list)), key=test_list.__getitem__)

# Printing result
print("Indices list of min K elements is : " + str(smallest_indices))

The original list is : [5, 6, 10, 4, 7, 1, 19]
Indices list of min K elements is : [5, 3, 0]

Time complexity: O(n log K) where n is the length of the list test list and K is the number of smallest elements required.
Auxiliary space: O(K), which is the size of the heap that is created by the heapq nsmallest() function to store the smallest K elements.

Method 3: Using numpy.argsort() function

The algorithm used in this approach is the argsort function of the NumPy library. The argsort function returns the indices that would sort the input array in ascending order. We pass the input list to the argsort function to obtain an array of indices in ascending order.

import numpy as np

# Input list
test_list = [5, 6, 10, 4, 7, 1, 19]

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

# Initialize K
K = 3

# Smallest K elements indices
# using numpy.argsort() function
res = np.argsort(test_list)[:K]

# Printing result
print("Indices list of min K elements is : " + str(list(res)))

The original list is : [5, 6, 10, 4, 7, 1, 19]
Indices list of min K elements is : [5, 3, 0]

Time complexity: O(n log n), where n is the length of the input list. This is because the argsort function uses the quicksort algorithm to sort the input list, which has an average time complexity of O(n log n). 
Auxiliary space: O(N), as we are creating a new array to store the sorted indices.

Method 4: Using the bisect module:

Steps:

1. Initialize a list test_list with some values and a variable K to a desired value.
2. Create a list of indices indices using the range function and the length of the test_list.
3. Sort the indices list based on the values in the test_list, using the sort method and a lambda function that takes an 4.index 'i' and returns the corresponding value in test_list[i].
4. Take the first K indices from the sorted list to get the smallest_indices.
5. Print the original list and the resulting smallest_indices list.

import bisect

# Input list and value
test_list = [5, 6, 10, 4, 7, 1, 19]
K = 3

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

# Creating a list of indices
indices = list(range(len(test_list)))

# Sorting indices based on values of test_list
indices.sort(key=lambda i: test_list[i])

# Getting the first K indices from the sorted list
smallest_indices = indices[:K]
# printing result
print("Indices list of min K elements is : " + str(smallest_indices))

# This code is contributed by Vinay pinjala

The original list is : [5, 6, 10, 4, 7, 1, 19]
Indices list of min K elements is : [5, 3, 0]

Time complexity: O(n log n), where n is the length of the input list test_list. This is because the algorithm involves sorting the list of indices using the sort() method, which has a time complexity of O(n log n) in the average and worst case.

Auxiliary space: O(n) because it creates a list of indices of length n and a list of K smallest indices, which has a maximum length of K. Therefore, the overall space complexity is dominated by the list of indices of length n.

Method 5 : Using extend(),sort() and index() methods

Approach

1. Create a new list with same elements of test_list using extend() method.
2. Sort that list using sort() function. 
3. Slice the list till K index. 
4. Initiate a for loop to traverse sliced list and append index of element of sliced list in original list to an output list.
5. Display output list.

# Python3 code to demonstrate working of
# Smallest K elements indices

# Initialize list
test_list = [5, 6, 10, 4, 7, 1, 19]

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

# Initialize K
K = 3

# Smallest K elements indices
res = []
x = []

x.extend(test_list)
x.sort()

for i in x[:K]:
    res.append(test_list.index(i))

    # printing result
print("Indices list of min K elements is : " + str(res))

The original list is : [5, 6, 10, 4, 7, 1, 19]
Indices list of min K elements is : [5, 3, 0]

Time complexity: O(n log n), where n is the length of the input list test_list. This is because the algorithm involves sorting the list of indices using the sort() method, which has a time complexity of O(n log n) in the average and worst case.
Auxiliary space: O(n) because it creates a list of indices of length n and a list of K smallest indices, which has a maximum length of K. Therefore, the overall space complexity is dominated by the list of indices of length n.

Method 6:  Using loop

Steps:

1. Initialize the list, test_list, and print it.
2. Initialize the value of K and print it.
3. Initialize an empty dictionary, index_dict, to store the indices and their corresponding values from test_list.
4. Use a for loop to iterate through test_list, and add each element and its index to index_dict.
5. Sort the dictionary index_dict by its values, which represent the elements of test_list.
6. Create a list, smallest_indices, to store the indices of the smallest K elements.
7. Use a for loop to iterate through the sorted index_dict and append the indices of the smallest K elements to the smallest_indices.
8. Print the list of indices of the smallest K elements.

# Initialize list
test_list = [5, 6, 10, 4, 7, 1, 19]

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

# Initialize K
K = 3

# Initializing dictionary to store indices and their corresponding values
index_dict = {}

# Adding each element and its index to the dictionary
for i, value in enumerate(test_list):
    index_dict[i] = value

# Sort the dictionary by its values (the elements of test_list)
sorted_dict = dict(sorted(index_dict.items(), key=lambda x: x[1]))

# Create a list to store the indices of the smallest K elements
smallest_indices = []

# Append the indices of the smallest K elements to the list
for i in range(K):
    smallest_indices.append(list(sorted_dict.keys())[i])

# Printing result
print("Indices list of min K elements is : " + str(smallest_indices))

The original list is : [5, 6, 10, 4, 7, 1, 19]
Indices list of min K elements is : [5, 3, 0]

Time complexity: O(nlogn), where n is the length of test_list.
Auxiliary space: O(n) to store the dictionary index_dict and the sorted dictionary sorted_dict

manjeet_04

Follow

Improve

Article Tags :

• Python
• Python Programs
• Python list-programs

Practice Tags :

• python