Group Sublists by another List using Python

In Python, we can group sublists by another list using various methods like dictionaries, itertools.groupby() function, and nested list comprehensions. Grouping sublists by another list is useful when analyzing large datasets, categorizing data, text analysis, and natural language processing. This article explores different methods to group sublists by another list in Python.

Using a Dictionary

A dictionary provides a straightforward approach to group sublists by another list in Python. This method creates groups based on a key and maintains the order specified by the grouping list.

Example

In this example, we define a function group_sublists that takes sublists and a grouping list as parameters. We create an empty dictionary to store sublists grouped by their keys ?

def group_sublists(sublists, grouping_list):
    groups = {}

    for sublist in sublists:
        key = sublist[0]  # Assuming the first element of each sublist is the key
        if key in groups:
            groups[key].append(sublist)
        else:
            groups[key] = [sublist]

    return [groups[key] for key in grouping_list]

# Example usage
sublists = [[1, 'apple'], [2, 'banana'], [1, 'orange'], [2, 'grape']]
grouping_list = [1, 2]

result = group_sublists(sublists, grouping_list)
print(result)

The output is ?

[[[1, 'apple'], [1, 'orange']], [[2, 'banana'], [2, 'grape']]]

Using itertools.groupby() Function

Python's itertools module provides the groupby() function that groups elements based on a key function. This method requires sorting the data first.

Syntax

itertools.groupby(iterable, key=None)

Here, iterable is the input sequence to group, and key is an optional function used for grouping. If no key function is provided, elements themselves are used as keys.

Example

We sort the sublists based on the key, then use itertools.groupby() to group them ?

import itertools

def group_sublists(sublists, grouping_list):
    sublists.sort(key=lambda x: x[0])  # Sort the sublists based on the key

    result = []
    for key, group in itertools.groupby(sublists, lambda x: x[0]):
        if key in grouping_list:
            result.append(list(group))

    return result

# Example usage
sublists = [[1, 'apple'], [2, 'banana'], [1, 'orange'], [2, 'grape']]
grouping_list = [1, 2]

result = group_sublists(sublists, grouping_list)
print(result)

The output is ?

[[[1, 'apple'], [1, 'orange']], [[2, 'banana'], [2, 'grape']]]

Using Nested List Comprehension

Nested list comprehension provides a concise way to group sublists by filtering elements based on matching keys.

Syntax

[expression for item in list if condition]

The syntax consists of square brackets enclosing an expression followed by a for loop. An optional if condition can filter elements.

Example

We use nested list comprehension to iterate over each key in the grouping list and filter matching sublists ?

def group_sublists(sublists, grouping_list):
    return [
        [sublist for sublist in sublists if sublist[0] == key]
        for key in grouping_list
    ]

# Example usage
sublists = [[1, 'apple'], [2, 'banana'], [1, 'orange'], [2, 'grape']]
grouping_list = [1, 2]

result = group_sublists(sublists, grouping_list)
print(result)

The output is ?

[[[1, 'apple'], [1, 'orange']], [[2, 'banana'], [2, 'grape']]]

Comparison

Conclusion

Use the dictionary method for most cases as it's efficient and preserves order. Choose itertools.groupby() for large datasets where memory efficiency matters. List comprehension works best for simple, readable solutions with small datasets.