Python – Group Elements in Matrix

Given a Matrix with two columns, group 2nd column elements on basis of 1st column.

Input : test_list = [[5, 8], [2, 0], [5, 4], [2, 3], [2, 9]] 
Output : {5: [8, 4], 2: [0, 3, 9]} 
Explanation : 8 and 4 are mapped to 5 in Matrix, all others to 2. 

Input : test_list = [[2, 8], [2, 0], [2, 4], [2, 3], [2, 9]] 
Output : {2: [8, 4, 0, 3, 9]} 
Explanation : All mapped to 2.

Method #1 : Using dictionary comprehension + loop

This is one of the ways in which this task can be performed. In this, we construct the dictionary with empty list values from row 1 and then run a loop to assign values into it.

Step-by-step approach:

• Initialize the list of lists as the input list.
• Print the original list.
• Initialize an empty dictionary with default empty lists using dictionary comprehension. The keys are the first elements of the tuples in the input list.
• Loop through the input list and append the second element of each tuple to the list in the dictionary with the same key.
• Print the result, which is the dictionary with the grouped elements.

Below is the implementation of the above approach:

Python3

# Python3 code to demonstrate working of
# Group Elements in Matrix
# Using dictionary comprehension + loop

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

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

# initializing empty dictionary with default empty list
res = {idx[0]: [] for idx in test_list}

# using loop for grouping
for idx in test_list:
    res[idx[0]].append(idx[1])

# printing result
print("The Grouped Matrix : " + str(res))

Output

The original list : [[5, 8], [2, 0], [5, 4], [2, 3], [7, 9]]
The Grouped Matrix : {5: [8, 4], 2: [0, 3], 7: [9]}

Time complexity: O(n), where n is the length of the input list.
Auxiliary space: O(n), as a dictionary is created with n key-value pairs.

Method #2 : Using loop + defaultdict()

This is similar to above method. The difference being that initial empty mesh is created using defaultdict().

Python3

# Python3 code to demonstrate working of
# Group Elements in Matrix
# Using loop + defaultdict()
from collections import defaultdict

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

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

# initializing empty dictionary using defaultdict
res = defaultdict(list)

# using loop for grouping
for idx in test_list:
    res[idx[0]].append(idx[1])

# printing result
print("The Grouped Matrix : " + str(dict(res)))

Output

The original list : [[5, 8], [2, 0], [5, 4], [2, 3], [7, 9]]
The Grouped Matrix : {5: [8, 4], 2: [0, 3], 7: [9]}

The time complexity of this code is O(n), where n is the number of elements in the input list. 

The auxiliary space complexity of this code is also O(n), where n is the number of elements in the input list.

Method #3 : Using nested for loops

Approach 

1. Extracting the unique 0 index values from nested list
2. Initiated a nested for loop to group similar first index elements

Python3

# Python3 code to demonstrate working of
# Group Elements in Matrix
# Using dictionary comprehension + loop

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

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

# initializing empty dictionary with default empty list
res=dict()
x=[]
for i in test_list:
    if i[0] not in x:
        x.append(i[0])
for i in x:
    v=[]
    for j in test_list:
        if j[0]==i:
            v.append(j[1])
    res[i]=v
# printing result
print("The Grouped Matrix : " + str(res))

Output

The original list : [[5, 8], [2, 0], [5, 4], [2, 3], [7, 9]]
The Grouped Matrix : {5: [8, 4], 2: [0, 3], 7: [9]}

Time Complexity : O(n*n)
Auxiliary Space : O(n)

Method #4: Using itertools.groupby()

Use the itertools.groupby() function to group elements in the matrix based on their first value. This function takes an iterable and a key function as arguments, and returns an iterator that produces pairs of keys and groups of items that have the same key.

Python3

from itertools import groupby

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

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

# grouping elements based on first value
test_list.sort()
groups = groupby(test_list, lambda x: x[0])

# creating dictionary from groups
res = {k: [v[1] for v in group] for k, group in groups}

# printing result
print("The Grouped Matrix : " + str(res))

Output

The original list : [[5, 8], [2, 0], [5, 4], [2, 3], [7, 9]]
The Grouped Matrix : {2: [0, 3], 5: [4, 8], 7: [9]}

Time complexity: O(n log n) due to the use of the groupby() function, which requires the input to be sorted. 
Auxiliary space: O(n), since we need to store the groups in a dictionary.

Method #5: Using list comprehension and set

This method uses list comprehension and set to group the elements in the matrix. The set function is used to get the unique elements of the first column of the matrix. Then, a dictionary comprehension is used to create a dictionary with the unique elements as keys and a list of the corresponding second column elements as values.

Python3

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

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

# using list comprehension and set to group elements
res = {i: [j[1] for j in test_list if j[0] == i] for i in set([k[0] for k in test_list])}

# printing result
print("The Grouped Matrix : " + str(res))

Output

The original list : [[5, 8], [2, 0], [5, 4], [2, 3], [7, 9]]
The Grouped Matrix : {2: [0, 3], 5: [8, 4], 7: [9]}

Time complexity: O(n^2), where n is the number of elements in the input list. 
Auxiliary space: O(n^2), as the resulting dictionary contains a key-value pair for each unique element in the first column of the input list.