如何根据 Python 中另一个列表的（子列表）索引对列表进行分区

Question

我有两个列表，一个包含一些唯一元素（在我的例子中是整数），另一个包含指示元素应该插入到新创建的嵌套列表的哪个子列表中的索引。

elements = [1, 2, 3, 4, 5, 6]
indices =  [0, 0, 1, 2, 2, 1]

expected_result = [[1, 2], [3, 6], [4, 5]]

元素列表仅包含唯一项，可能未排序。 索引列表是“标准化”的，这样较低的索引总是首先出现。 新的嵌套列表应使用索引来确定元素应属于的预期结果的子列表。

我想出了下面这个函数，但是我感觉应该有一个更简单的方法。

def indices_to_nested_lists(indices: Sequence[int], elements: Sequence):
    result = []
    for i in range(max(indices)+1):
        sublist = []
        for j in range(len(elements)):
            if indices[j] == i:
                sublist.append(elements[j])
        result.append(sublist)
    return result

谁能想到一种更简单、可能更 Python 的方式来实现相同的结果？

Answer 1

尝试将此 for 循环与 zip 一起使用：

l = [[] for i in range(max(indices) + 1)]
for x, y in zip(elements, indices):
    l[y].append(x)
print(l)

输出：

[[1, 2], [3, 6], [4, 5]]