How to partition a list based on (sublist) indices of another list in Python

Question

I have two lists, one containing some unique elements (integers in my case) and the other containing indices that indicate into which sublist of a newly created nested list the elements should be inserted.

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

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

The list of elements contains only unique items, potentially not sorted. The list of indices is 'normalized' such that the lower indices will always occur first. The new nested list should use the indices to determine the sublist of the expected result to which the elements shall belong.

I have come up with the following function, but I have a feeling that there should be an easier way.

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

Can anyone think of an easier, maybe more pythonic way of achieving the same result?

U13-Forward · Accepted Answer · 2021-07-15T12:29:03.023

2

Try using this for loop with zip:

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

Output:

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

edited Jul 15 '21 at 12:29

answered Jul 15 '21 at 07:33

U13-Forward

69,221
14
89
114

```TypeError: can only concatenate list (not "int") to list```. I get an error – Jul 15 '21 at 07:42
It's just a typo — it should be `l = [[] for i in range(max(indices) + 1)]`. It's a great answer. You could also write that line as `l = [[]] * (max(indices) + 1)`. – Alex Waygood Jul 15 '21 at 08:53
Ah, yes so much nicer! Thanks :) – eltings Jul 15 '21 at 13:28

How to partition a list based on (sublist) indices of another list in Python

1 Answers1