Group the People Given the Group Size They Belong To
Problem
There are n people that are split into some unknown number of groups. Each person is labeled with a unique ID from 0 to n - 1.
You are given an integer array groupSizes, where groupSizes[i] is the size of the group that person i is in. For example, if groupSizes[1] = 3, then person 1 must be in a group of size 3.
Return a list of groups such that each person i is in a group of size
groupSizes[i].
Each person should appear in exactly one group , and every person must be in a group. If there are multiple answers, return any of them. It is guaranteed that there will be at least one valid solution for the given input.
Examples
Example 1
Input: groupSizes = [3,3,3,3,3,1,3]
Output: [[5],[0,1,2],[3,4,6]]
Explanation:
The first group is [5]. The size is 1, and groupSizes[5] = 1.
The second group is [0,1,2]. The size is 3, and groupSizes[0] = groupSizes[1] = groupSizes[2] = 3.
The third group is [3,4,6]. The size is 3, and groupSizes[3] = groupSizes[4] = groupSizes[6] = 3.
Other possible solutions are [[2,1,6],[5],[0,4,3]] and [[5],[0,6,2],[4,3,1]].
Example 2
Input: groupSizes = [2,1,3,3,3,2]
Output: [[1],[0,5],[2,3,4]]
Constraints
groupSizes.length == n1 <= n <= 5001 <= groupSizes[i] <= n
Solution
Method 1 – Hash Map Grouping and Greedy Assignment
Intuition
We can use a hash map to collect people by their required group size. As soon as we collect enough people for a group, we finalize that group and start a new one. This greedy approach ensures all constraints are satisfied.
Approach
- Initialize a hash map to store lists of people for each group size.
- Iterate through
groupSizes, for each personi:- Add
ito the list forgroupSizes[i]. - If the list reaches the required group size, add it to the answer and reset the list.
- Add
- Return the list of groups.
Code
C++
class Solution {
public:
vector<vector<int>> groupThePeople(vector<int>& groupSizes) {
unordered_map<int, vector<int>> mp;
vector<vector<int>> ans;
for (int i = 0; i < groupSizes.size(); ++i) {
mp[groupSizes[i]].push_back(i);
if (mp[groupSizes[i]].size() == groupSizes[i]) {
ans.push_back(mp[groupSizes[i]]);
mp[groupSizes[i]].clear();
}
}
return ans;
}
};
Go
func groupThePeople(groupSizes []int) [][]int {
mp := map[int][]int{}
var ans [][]int
for i, sz := range groupSizes {
mp[sz] = append(mp[sz], i)
if len(mp[sz]) == sz {
ans = append(ans, append([]int{}, mp[sz]...))
mp[sz] = nil
}
}
return ans
}
Java
class Solution {
public List<List<Integer>> groupThePeople(int[] groupSizes) {
Map<Integer, List<Integer>> mp = new HashMap<>();
List<List<Integer>> ans = new ArrayList<>();
for (int i = 0; i < groupSizes.length; i++) {
int sz = groupSizes[i];
mp.computeIfAbsent(sz, k -> new ArrayList<>()).add(i);
if (mp.get(sz).size() == sz) {
ans.add(new ArrayList<>(mp.get(sz)));
mp.get(sz).clear();
}
}
return ans;
}
}
Kotlin
class Solution {
fun groupThePeople(groupSizes: IntArray): List<List<Int>> {
val mp = mutableMapOf<Int, MutableList<Int>>()
val ans = mutableListOf<List<Int>>()
for (i in groupSizes.indices) {
val sz = groupSizes[i]
mp.getOrPut(sz) { mutableListOf() }.add(i)
if (mp[sz]!!.size == sz) {
ans.add(mp[sz]!!.toList())
mp[sz]!!.clear()
}
}
return ans
}
}
Python
class Solution:
def groupThePeople(self, groupSizes: list[int]) -> list[list[int]]:
from collections import defaultdict
mp: dict[int, list[int]] = defaultdict(list)
ans: list[list[int]] = []
for i, sz in enumerate(groupSizes):
mp[sz].append(i)
if len(mp[sz]) == sz:
ans.append(mp[sz][:])
mp[sz].clear()
return ans
Rust
impl Solution {
pub fn group_the_people(group_sizes: Vec<i32>) -> Vec<Vec<i32>> {
use std::collections::HashMap;
let mut mp: HashMap<i32, Vec<i32>> = HashMap::new();
let mut ans = Vec::new();
for (i, &sz) in group_sizes.iter().enumerate() {
let e = mp.entry(sz).or_insert(Vec::new());
e.push(i as i32);
if e.len() == sz as usize {
ans.push(e.clone());
e.clear();
}
}
ans
}
}
TypeScript
class Solution {
groupThePeople(groupSizes: number[]): number[][] {
const mp: Record<number, number[]> = {};
const ans: number[][] = [];
for (let i = 0; i < groupSizes.length; i++) {
const sz = groupSizes[i];
if (!mp[sz]) mp[sz] = [];
mp[sz].push(i);
if (mp[sz].length === sz) {
ans.push([...mp[sz]]);
mp[sz] = [];
}
}
return ans;
}
}
Complexity
- ⏰ Time complexity:
O(n)— Each person is processed once. - 🧺 Space complexity:
O(n)— For storing group assignments.