Python按N个块拆分列表 拆分点不拆分同首位编号的字符串组

需求为将字符串列表拆分为N个大小尽可能均等的块，拆分需遵守强制规则：

• 拆分点仅允许出现在字符串首位编号切换的位置，同属一个首位编号的元素必须在同一个块内，不可拆分
• 字符串固定为[任意前缀]-X-Y的两段数字后缀格式，取第一段数字X作为分组标识，X、Y支持多位数，例如file-123-45属于编号123的分组

原有实现的问题 #

原有代码仅按列表长度做均等切片，完全不考虑编号分组规则，会把同属一个首位编号的元素拆分到不同块中。
原有代码如下：

def split(a, n):
    k, m = divmod(len(a), n)
    return (a[i*k+min(i, m):(i+1)*k+min(i+1, m)] for i in range(n))


l = ["string-1-1", "string-1-2", "string-2-1", "string-2-2", "string-2-3"]
chunks = 2
l_split = list(split(l, chunks))

原有代码运行输出不符合规则：

[['string-1-1', 'string-1-2', 'string-2-1'], ['string-2-2', 'string-2-3']]

上述结果中编号为2的三个元素被拆分到两个块中，违反规则。

注：提问中给出的预期输出最后一个元素-2-3为笔误，对应原列表中最后一个元素string-2-3，正确预期输出为编号1的两个元素为一个块，编号2的三个元素为一个块。

修改后实现 #

实现分两步：

1. 先遍历原列表，按首位编号把列表拆分为不可再切割的原子组，同组元素编号完全一致，不允许拆分
2. 采用贪心策略把原子组合并为N个块，每次判断当前块累计长度和平均目标长度的差值，决定是把当前原子组加入当前块还是新建块，保证最终块大小尽可能均等

完整代码：

def split(a, n):
    if n <= 0:
        raise ValueError("拆分块数必须为正整数")
    if len(a) == 0:
        return [[] for _ in range(n)]
    
    # 第一步：提取首位编号，拆分不可切割的原子组
    groups = []
    # 取第一个元素的首位编号
    current_prefix = int(a[0].split('-')[-2])
    current_group = [a[0]]
    for item in a[1:]:
        prefix = int(item.split('-')[-2])
        if prefix == current_prefix:
            current_group.append(item)
        else:
            groups.append(current_group)
            current_group = [item]
            current_prefix = prefix
    groups.append(current_group)

    # 原子组数量不足n时，补空块满足返回n块的要求
    if len(groups) < n:
        return groups + [[] for _ in range(n - len(groups))]

    # 第二步：贪心合并原子组，保证块大小尽可能平均
    total_length = len(a)
    target_avg = total_length / n
    chunks = []
    current_chunk = []
    current_len = 0
    # 预留最后一块的位置，避免块数超过n
    for idx, group in enumerate(groups):
        group_len = len(group)
        # 当前块非空，且加完当前组会超过平均大小、或者已经到最后n-1块的位置，就把当前块存入结果
        if current_chunk and (current_len + group_len / 2 > target_avg or len(chunks) >= n - 1):
            chunks.append(current_chunk)
            current_chunk = []
            current_len = 0
        current_chunk.extend(group)
        current_len += group_len
    # 加入最后一个块
    if current_chunk:
        chunks.append(current_chunk)
    
    return chunks

# 测试
l = ["string-1-1", "string-1-2", "string-2-1", "string-2-2", "string-2-3"]
chunks = 2
l_split = split(l, chunks)
print(l_split)

运行结果 #

[['string-1-1', 'string-1-2'], ['string-2-1', 'string-2-2', 'string-2-3']]

完全符合拆分规则，且块大小尽可能均等。

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