DEV Community

Cover image for Algorithms Problem Solving: Group the people
TK
TK

Posted on • Originally published at leandrotk.github.io

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]]
Enter fullscreen mode Exit fullscreen mode

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]
Enter fullscreen mode Exit fullscreen mode

...into this hash map:

{
    1: [2],
    2: [3, 5],
    3: [0, 1, 4, 6, 7, 8]
}
Enter fullscreen mode Exit fullscreen mode

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]
Enter fullscreen mode Exit fullscreen mode

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]
}
Enter fullscreen mode Exit fullscreen mode

To this:

[[2], [3, 5], [0, 1, 4], [6, 7, 8]]
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

Resources

Discussion (1)

Collapse
musale profile image
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.