Python | Indices of sorted list of list elements

Sorting is a common construct and there have been many variations of it being discussed. But sometimes, we need to perform the sorting on the list of list and moreover just require to find the order in which element occurs before getting sorted. Let's find out how to get indices of sorted order in list of lists.

Method #1 : Using List comprehension + enumerate() + sort() The combination of above 3 functions can be used to perform this particular task. In this, we perform the sorting taking triplets consisting of element and row and column coordinates and return them in 2nd step. 

# Python3 code to demonstrate
# Indices of sorted list of list elements
# using List comprehension + enumerate() + sort()

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

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

# using List comprehension + enumerate() + sort()
# Indices of sorted list of list elements
res = [(i, j) for i, x in enumerate(test_list)
                     for j, k in enumerate(x)]

res.sort(key = lambda ij: test_list[ij[0]][ij[1]]) 

# print result
print("The indices of sorted order are : " + str(res))

Time complexity: O(n^2 log n), where n is the total number of elements in the list of lists. 
Auxiliary space: O(n), where n is the total number of elements in the list of lists. 

Method #2 : Using sorted() + lambda The task performed above can be performed as arguments to the sorted function and lambda function performs the task of list comprehension function as above. 

# Python3 code to demonstrate
# Indices of sorted list of list elements
# using sorted() + lambda

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

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

# using sorted() + lambda
# Indices of sorted list of list elements
res = sorted([(i, j) for i, x in enumerate(test_list)
                            for j, k in enumerate(x)],
             key = lambda ij: test_list[ij[0]][ij[1]])

# print result
print("The indices of sorted order are : " + str(res))

Time complexity: O(n^2 * log n).  where N is the number of sublists in the test_list and M is the maximum length of any sublist
Auxiliary space: O(n^2).where N is the number of sublists in the input list and M is the maximum length of any sublist. 

Method #3: Using heapq.nsmallest().

The algorithm used in the given code is as follows:

1. Initialize an empty list res.
2. For each row in test_list, find the smallest len(x) elements using heapq.nsmallest().
3. For each element found in step 2, append a tuple (i, x.index(k)) to res, where i is the row index and x.index(k) is the column index of the element k.
4. Sort the list of tuples res by the values of test_list at the corresponding indices.
5. Print the sorted list of tuples.

import heapq
# initializing list
test_list = [[4, 5, 1],
             [9, 3, 2],
             [8, 6]]

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


res = [(i, x.index(k)) for i, x in enumerate(test_list) for k in heapq.nsmallest(len(x), x)]
res.sort(key=lambda ij: test_list[ij[0]][ij[1]])
print("The indices of sorted order are:", res)

The original list : [[4, 5, 1], [9, 3, 2], [8, 6]]
The indices of sorted order are: [(0, 2), (1, 2), (1, 1), (0, 0), (0, 1), (2, 1), (2, 0), (1, 0)]

The time complexity of the algorithm is O(N*M*log(M)), where N is the number of rows in test_list and M is the maximum number of columns in any row. This is because for each row, we use heapq.nsmallest() to find the smallest M elements, which takes O(M*log(M)) time. Since we do this N times, the overall time complexity is O(N*M*log(M)).

The auxiliary space of the algorithm is O(N*M), since we store the list of tuples res that contains one tuple for each element in test_list.

Method 4 :  Using Flattening and Sorting

1. Initialize a list of lists test_list with integer values.
2. Print the original list using the print() function.
3. Use list comprehension to create a flattened list of all elements in the test_list and assign it to the variable flat_list. Also, create a list of indices from 0 to the length of flat_list - 1 and assign it to the variable indices.
4. Sort the indices list based on the corresponding element in the flat_list. This is done using the sort() method with a key parameter that takes a lambda function that returns the value at the corresponding index of flat_list.
5. Calculate the row and column indices from the sorted indices list by dividing each index by the number of columns in the original test_list and getting the floor division quotient as rows. The remainder after division gives the column index which is stored in cols.
6. Zip the rows and cols lists together to create a list of tuples that represent the indices of the sorted order and assign it to the variable res.
7. Print the result using the print() function.

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

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

# get flattened list and corresponding indices
flat_list = [element for sublist in test_list for element in sublist]
indices = list(range(len(flat_list)))

# sort the indices based on the corresponding element in flat_list
indices.sort(key=lambda x: flat_list[x])

# get row and column indices from sorted indices
rows = [index // len(test_list[0]) for index in indices]
cols = [index % len(test_list[0]) for index in indices]

# zip row and column indices and convert to list of tuples
res = list(zip(rows, cols))

# print result
print("The indices of sorted order are : " + str(res))

The original list : [[4, 5, 1], [9, 3, 2], [8, 6]]
The indices of sorted order are : [(0, 2), (1, 2), (1, 1), (0, 0), (0, 1), (2, 1), (2, 0), (1, 0)]

The time complexity of this method is O(nlogn), where n is the total number of elements in the nested list, because of the sorting operation. 

The auxiliary space is also O(n) because we create a separate list to store the flattened version of the nested list and another list to store the corresponding indices.