如何在Python中基于索引子列表对元素列表分组/分批？

需求规则 #

需要实现一个分组函数，遵循以下逻辑处理elements列表：

• 优先按照indices中的元素依次从elements中提取对应数量的元素作为分组；
• 若elements总长度大于indices元素总和：剩余元素以indices中的最大值为每组大小进行分组，直到元素耗尽；
• 若elements总长度小于indices元素总和：按照indices的顺序取元素，直到elements耗尽，最后一组取剩余的所有元素。

实现代码 #

def group_elements(indices, elements):
    result = []
    current_pos = 0
    max_batch_size = max(indices) if indices else 0

    # 先处理indices定义的初始分组
    for count in indices:
        if current_pos >= len(elements):
            break
        # 计算当前分组的结束位置，不超过元素列表长度
        end_pos = current_pos + count
        batch = elements[current_pos:end_pos]
        if batch:
            result.append(batch)
        current_pos = end_pos

    # 处理剩余未分组的元素
    while current_pos < len(elements):
        end_pos = current_pos + max_batch_size
        result.append(elements[current_pos:end_pos])
        current_pos = end_pos

    return result

场景验证 #

场景1：elements长度等于indices元素总和 #

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I']
print(group_elements(indices, elements))
# 输出: [['A'], ['B', 'C', 'D'], ['E', 'F', 'G', 'H', 'I']]

场景2：elements长度大于indices总和（剩余元素刚好凑一组最大值） #

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N']
print(group_elements(indices, elements))
# 输出: [['A'], ['B', 'C', 'D'], ['E', 'F', 'G', 'H', 'I'], ['J', 'K', 'L', 'M', 'N']]

场景3：elements长度大于indices总和（剩余元素不足一组最大值） #

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L']
print(group_elements(indices, elements))
# 输出: [['A'], ['B', 'C', 'D'], ['E', 'F', 'G', 'H', 'I'], ['J', 'K', 'L']]

场景4：elements长度大于indices总和（剩余元素可分多组最大值） #

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F','G','H','I', 'J','K','L','M','N','O','P','Q','R','S']
print(group_elements(indices, elements))
# 输出: [['A'], ['B', 'C', 'D'], ['E', 'F', 'G', 'H', 'I'], ['J', 'K', 'L', 'M', 'N'], ['O', 'P', 'Q', 'R', 'S']]

场景5：elements长度小于indices总和（未完成前两个分组） #

indices = [1, 3, 5]
elements = ['A','B','C']
print(group_elements(indices, elements))
# 输出: [['A'], ['B', 'C']]

场景6：elements长度小于indices总和（完成前两个分组，第三个分组取剩余） #

indices = [1, 3, 5]
elements = ['A','B','C','D','E','F']
print(group_elements(indices, elements))
# 输出: [['A'], ['B', 'C', 'D'], ['E', 'F']]

内容的提问来源于stack exchange，提问作者Learner