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
Top comments (1)
I could use a nifty one-line to convert
grouped_by_size
dict togrouped_by_ids
list like this:[*grouped_by_size.values()]
. It's in the order of thegrouped_by_size
keys.