Python | Group elements at same indices in a multi-list

Flattening a 2D list to one is a common problem that is faced in many domains. But sometimes we require to pair elements at specific indices as one, so that elements at respective indices are together. This problem is not common but still having a solution to it helps. Let’s discuss certain ways to pair elements at specific indices. 

Method #1 : Using list comprehension + zip() List comprehension can be used to achieve this particular task along with zip function which does the task of pairing the like indices together. This method is just shorthand to the naive method. 

Python3

# Python3 code to demonstrate
# index list elements pairing
# using list comprehension

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

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

# using list comprehension
# to perform index list elements pairing
res = [list(x) for x in zip(*test_list)]

# printing result
print("The index elements pairs list is " + str(res))

Output:

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time Complexity: O(n), where n is the number of elements in the list.
Auxiliary Space: O(n), as a new list ‘res’ is created which contains the padded elements.

Method #2: Using map() + zip() map function can be used to map each iteration result into a single list while zip function performs the index element pairing. This combination can be used to achieve the desired result.

Python3

# Python3 code to demonstrate
# index list elements pairing
# using map() + zip()

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

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

# using map() + zip()
# to perform index list elements pairing
res = list(map(list, zip(*test_list)))

# printing result
print("The index elements pairs list is " + str(res))

Output:

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time complexity: O(n^2) where n is the length of the longest list in test_list. 
Auxiliary space: O(n) where n is the length of the longest list in test_list.

Method #3: Using zip()

Using zip() function alone does it all and just needs to be typecast to list to print result in list format. It’s the most pythonic way to perform this task and the most elegant as well. 

Python3

# Python3 code to demonstrate
# index list elements pairing
# using zip()

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

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

# using zip()
# to perform index list elements pairing
res = list(zip(*test_list))

# printing result
print ("The index elements pairs list is " + str(res))

Output:

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [(1, 4, 8), (4, 6, 3), (5, 8, 10)]

Time complexity: O(nm), where n is the number of lists and m is the length of each list.
Auxiliary space: O(nm), as the zip() function returns a new list with n elements, each containing m elements. This list is stored in the res variable, which takes up O(nm) space.

Method #4 : Using numpy():

1. Import the NumPy library using import numpy as np.
2. Initialize a list of lists test_list with some values.
3. Print the original list using print(“The original list is : ” + str(test_list)).
4. Use NumPy to transpose the list of lists and convert the result back to a list using np.array(test_list).T.tolist().
5. Print the result using print(“The index elements pairs list is ” + str(res)).

Python3

# Python3 code to demonstrate
# index list elements pairing
# using NumPy
import numpy as np

# initializing input list
test_list = [[1, 4, 5], [4, 6, 8], [8, 3, 10]]

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

# using NumPy # to perform index list elements pairing
res = np.array(test_list).T.tolist()

# printing result
print("The index elements pairs list is " + str(res))

# This code is contributed by Jyothi pinjala.

Output:

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time complexity: O(N^2), where N is the size of the input list. This is because the algorithm involves iterating over each element of each list in the input list.

Auxiliary space: O(N^2), because the NumPy array function is used to create a new array that is the transpose of the input list of lists, which has the same number of elements as the original list. The tolist function is then called on this array to convert it back to a list, which also has the same number of elements as the original list. Therefore, the space used by the new list is proportional to the size of the original list.

Method #5: Using the itertools.zip_longest() function.

Algorithm:

1. Import the itertools module.
2. Initialize the original list test_list with nested sublists.
3. Print the original list.
4. Use the zip_longest() function from the itertools module to pair the elements in each sublist by index, padding the shorter sublists with None values.
5. Use a list comprehension to filter out the None values in each paired sublist and store the result in the paired_list variable.
6. Print the paired list.

Python3

import itertools

# Initializing list
test_list = [[1, 4, 5], [4, 6, 8], [8, 3, 10]]

# printing original list
print("The original list is : " + str(test_list))
paired_list = [[x for x in sublist if x is not None]
               for sublist in itertools.zip_longest(*test_list)]

# printing result
print("The index elements pairs list is " + str(paired_list))

# This code is contributed by tvsk

Output

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time complexity: O(n * m), where n is the length of the sublists in the original list and m is the length of the original list. This is because the zip_longest() function iterates over each element in each sublist, which requires n * m iterations in total.
Auxiliary Space: O(n * m), where n is the length of the sublists in the original list and m is the length of the original list. This is because the paired list requires a new list to be created with n * m elements, and any None values in the paired list must also be stored in memory.

Method #6: Using a for loop and a temporary list.

Approach:

1. Initialize an empty list res.
2. Loop through the range of the length of any sublist in test_list.
3. For each index, loop through the sublists in test_list and append the element at that index to a temporary list temp.
4. Append temp to res.
5. Return res.

Below is the implementation of the above approach:

Python3

# Python3 code to demonstrate
# index list elements pairing
# using a for loop and a temporary list

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

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

# using a for loop and a temporary list
# to perform index list elements pairing
res = []
for i in range(len(test_list[0])):
    temp = []
    for sublist in test_list:
        temp.append(sublist[i])
    res.append(temp)

# printing result
print("The index elements pairs list is " + str(res))

Output

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time complexity: O(nm), where n is the number of sublists in test_list and m is the length of the sublists.
Auxiliary space: O(nm), for the res list.

Method #7: Using a list comprehension and nested loops

Steps:

1. Initialize the list test_list with the given values.
2. Print the original list using print() function.
3. Use a list comprehension and nested loops to perform index list elements pairing. In this step, for each element i in the range of the length of the first sublist of the test_list, create a temporary list temp containing all the i-th elements of each sublist. Then append this temp list to the res list.
4. Print the result using print() function.

Python3

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

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

# using a list comprehension and nested loops
# to perform index list elements pairing
res = [[sublist[i] for sublist in test_list] for i in range(len(test_list[0]))]

# printing result
print("The index elements pairs list is " + str(res))

Output

The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time complexity: O(n^2), where n is the number of sublists in test_list.
Auxiliary space: O(n^2), where n is the number of sublists in test_list.

Method #8: Using Pandas Dataframe

1. Import the Pandas library
2. Create a DataFrame from the input list
3. Use the transpose() method to swap the rows and columns of the DataFrame
4. Convert the transposed DataFrame to a list
5. Print the output list

Python3

import pandas as pd

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

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

# creating a dataframe from the input list
df = pd.DataFrame(test_list)

# using transpose method to swap the rows and columns of the dataframe
df_transposed = df.transpose()

# converting the transposed dataframe to a list
res = df_transposed.values.tolist()

# printing result
print("The index elements pairs list is " + str(res))

OUTPUT : 
The original list is : [[1, 4, 5], [4, 6, 8], [8, 3, 10]]
The index elements pairs list is [[1, 4, 8], [4, 6, 3], [5, 8, 10]]

Time complexity: O(n^2) (due to iterating over all elements of the input list)
Auxiliary space: O(n^2) (due to storing all the elements in the output list)