参考:

https://leetcode-cn.com/problems/group-the-people-given-the-group-size-they-belong-to/solution/yong-hu-fen-zu-by-leetcode-solution/

from typing import List


class Solution:
    def groupThePeople(self, groupSizes: List[int]) -> List[List[int]]:
        l = list()
        d = dict()
        # 先按每个元素分组
        for i in range(len(groupSizes)):
            if groupSizes[i] in d:
                d[groupSizes[i]].append(i)
            else:
                d[groupSizes[i]] = [i, ]
        # 每组元素按元素值拆分
        for k, v in d.items():
            start = 0
            while len(v) > start:
                l.append(v[start:start + k])
                start += k
        return l


if __name__ == '__main__':
    x = Solution()
    groupSizes = [3, 3, 3, 3, 3, 1, 3]
    print(x.groupThePeople(groupSizes))
    groupSizes = [2, 1, 3, 3, 3, 2]
    print(x.groupThePeople(groupSizes))