如何获取二维列表按指定规则分组排序后的分区索引

我有一个二维列表，需要按嵌套列表的第一个值对原列表的索引进行分组，再根据嵌套列表的第二个值对每个分组内的索引进行排序，最终返回排序后的分组索引列表。示例如下：

test = [[1, 2], [1, 1], [1, 5], [2, 3], [2, 1], [1, 10]]
sorted_partitions(test)
>>> [[1, 0, 2, 5], [4, 3]]
# 分组依据为 [(1, [1, 1]), (0, [1, 2]), (2, [1, 5]), (5, [1, 10]), (4, [2, 1]), (3, [2, 3])]

方法一：字典分组+自定义排序 #

先通过字典将索引按子列表的第一个元素归类，再对每个分组内的索引，依据原列表对应子项的第二个值排序：

def sorted_partitions(lst):
    groups = {}
    # 遍历列表，按子列表第一个元素分组索引
    for idx, item in enumerate(lst):
        key = item[0]
        if key not in groups:
            groups[key] = []
        groups[key].append(idx)
    # 对每个分组的索引按对应子列表第二个值排序
    return [sorted(indices, key=lambda i: lst[i][1]) for indices in groups.values()]

# 测试验证
test = [[1, 2], [1, 1], [1, 5], [2, 3], [2, 1], [1, 10]]
print(sorted_partitions(test))  # 输出: [[1, 0, 2, 5], [4, 3]]

方法二：使用itertools.groupby（需预排序） #

借助itertools.groupby分组，但要先将带索引的元素按分组键排序，确保同组元素连续：

from itertools import groupby

def sorted_partitions(lst):
    # 生成包含分组键、排序值、索引的元组列表
    indexed_items = [(item[0], item[1], idx) for idx, item in enumerate(lst)]
    # 先按分组键排序，保证groupby能正确分组
    indexed_items.sort(key=lambda x: x[0])
    # 分组后对每组索引按排序值排序
    sorted_groups = []
    for _, group in groupby(indexed_items, key=lambda x: x[0]):
        sorted_indices = sorted((g[2] for g in group), key=lambda i: lst[i][1])
        sorted_groups.append(sorted_indices)
    return sorted_groups

# 测试验证
test = [[1, 2], [1, 1], [1, 5], [2, 3], [2, 1], [1, 10]]
print(sorted_partitions(test))  # 输出: [[1, 0, 2, 5], [4, 3]]

核心逻辑说明 #

1. 分组：将原列表的索引按照子列表第一个元素归类，比如所有第一个元素为1的索引归为一组，2的归为另一组。
2. 排序：对每个分组内的索引，根据原列表中该索引对应子列表的第二个元素值做升序排序，最终得到符合要求的分组索引列表。

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