Python- Find the indices for k Smallest Elements

Finding indices of k smallest elements is a common task in data analysis and algorithm problems. Python provides several efficient approaches using built-in functions like sorted(), heapq, and NumPy's argsort().

Using sorted() with Lambda Function

This approach sorts indices based on their corresponding values ?

def k_smallest_indices(numbers, k):
    sorted_indices = sorted(range(len(numbers)), key=lambda i: numbers[i])
    return sorted_indices[:k]

# Create the list
numbers = [50, 20, 90, 10, 70]
k = 3

small_indices = k_smallest_indices(numbers, k)
print("The k smallest indices are:", small_indices)
print("Values at those indices:", [numbers[i] for i in small_indices])

The k smallest indices are: [3, 1, 0]
Values at those indices: [10, 20, 50]

Using heapq.nsmallest()

The heap-based approach is efficient for finding k smallest elements ?

import heapq

def find_k_smallest_indices(numbers, k):
    # Get k smallest values with their indices
    indexed_values = [(val, i) for i, val in enumerate(numbers)]
    smallest_pairs = heapq.nsmallest(k, indexed_values)
    return [index for value, index in smallest_pairs]

# Create the list
numbers = [5, 4, 3, 2, 1]
k = 3

smallest_indices = find_k_smallest_indices(numbers, k)
print("The k smallest indices are:", smallest_indices)
print("Values at those indices:", [numbers[i] for i in smallest_indices])

The k smallest indices are: [4, 3, 2]
Values at those indices: [1, 2, 3]

Using numpy.argsort()

NumPy's argsort() returns indices that would sort the array ?

import numpy as np

def find_k_smallest_indices(numbers, k):
    indices = np.argsort(numbers)
    return indices[:k].tolist()

# Create the list
numbers = [55, 7, 12, 2, 67, 90, 9]
k = 3

smallest_indices = find_k_smallest_indices(numbers, k)
print("The k smallest indices are:", smallest_indices)
print("Values at those indices:", [numbers[i] for i in smallest_indices])

The k smallest indices are: [3, 1, 6]
Values at those indices: [2, 7, 9]

Using Dictionary for Duplicate Values

This approach handles duplicate values by storing all indices for each value ?

def find_k_smallest_indices(numbers, k):
    # Group indices by value
    value_to_indices = {}
    for i, val in enumerate(numbers):
        if val not in value_to_indices:
            value_to_indices[val] = []
        value_to_indices[val].append(i)
    
    # Get indices for k smallest values
    result_indices = []
    for val in sorted(set(numbers)):
        result_indices.extend(value_to_indices[val])
        if len(result_indices) >= k:
            break
    
    return result_indices[:k]

# Create the list with duplicates
numbers = [5, 2, 9, 1, 2, 7]
k = 4

smallest_indices = find_k_smallest_indices(numbers, k)
print("The k smallest indices are:", smallest_indices)
print("Values at those indices:", [numbers[i] for i in smallest_indices])

The k smallest indices are: [3, 1, 4, 0]
Values at those indices: [1, 2, 2, 5]

Comparison

Conclusion

Use heapq.nsmallest() for optimal performance when k is small compared to n. For general cases, sorted() with lambda provides a clean and readable solution. NumPy's argsort() is ideal when working with numerical arrays.