Algorithms Problem Solving: Group the people

This post is part of the Algorithms Problem Solving series.

Problem description

This is the Group the People problem. The description looks like this:

There are n people whose IDs go from 0 to n - 1 and each person belongs exactly to one group. Given the array groupSizes of length n telling the group size each person belongs to, return the groups there are and the people's IDs each group includes.

You can return any solution in any order and the same applies for IDs. Also, it is guaranteed that there exists at least one solution.

Examples

Input: groupSizes = [3,3,3,3,3,1,3]
Output: [[5],[0,1,2],[3,4,6]]

Input: groupSizes = [2,1,3,3,3,2]
Output: [[1],[0,5],[2,3,4]]

Solution

My first idea was to just group the sizes into a hash map where the key is the size and the value is a list of all the indices for a specific size. An example would be to transform this list:

[3, 3, 1, 2, 3, 2, 3, 3, 3]

...into this hash map:

{
    1: [2],
    2: [3, 5],
    3: [0, 1, 4, 6, 7, 8]
}

The code is very simple. You start an empty hash map and then populate while iterating the list:

grouped_by_size = {}

for index in range(len(group_sizes)):
    size = group_sizes[index]

    if size in grouped_by_size:
        grouped_by_size[size] += [index]
    else:
        grouped_by_size[size] = [index]

With this organized data, we just need to group them by the id.

We will go from this:

{
    1: [2],
    2: [3, 5],
    3: [0, 1, 4, 6, 7, 8]
}

To this:

[[2], [3, 5], [0, 1, 4], [6, 7, 8]]

The idea here is to create an empty list to add all indices. Then iterate through the hash map and for each key-value, we add the value to the list. But we need to consider that our lists can be greater the maximum permitted length. So we cut them by the maximum length.

grouped_by_ids = []

for size, indices in grouped_by_size.items():
    for index in range(0, len(indices), size):
        grouped_by_ids.append(indices[index:index+size])

return grouped_by_ids

Now we have the grouped_by_ids list with all indices.

The whole algorithm looks like this:

def group_the_people(group_sizes):
    grouped_by_size = {}

    for index in range(len(group_sizes)):
        size = group_sizes[index]

        if size in grouped_by_size:
            grouped_by_size[size] += [index]
        else:
            grouped_by_size[size] = [index]

    grouped_by_ids = []

    for size, indices in grouped_by_size.items():
        for index in range(0, len(indices), size):
            grouped_by_ids.append(indices[index:index+size])

    return grouped_by_ids

Resources

• Learning Python: From Zero to Hero
• Algorithms Problem Solving Series
• Big-O Notation For Coding Interviews and Beyond
• Learn Python from Scratch
• Learn Object-Oriented Programming in Python
• Data Structures in Python: An Interview Refresher
• Data Structures and Algorithms in Python
• Data Structures for Coding Interviews in Python
• One Month Python Course

Comments

[Musale Martin]

I could use a nifty one-line to convert grouped_by_size dict to grouped_by_ids list like this: [*grouped_by_size.values()]. It's in the order of the grouped_by_size keys.