Blackhammer
Blackhammer

Reputation: 45

Partition a list items using a list containing partition numbers

Given a list

lst=['a','a','b','b','b','c','d','d']

and a list 'l' containing partition numbers

l=[2,3,1,2]

what I want is

partitioned_lst=[['a','a'],['b','b','b'],['c'],['d','d']]

Upvotes: 1

Views: 260

Answers (6)

juanpa.arrivillaga
juanpa.arrivillaga

Reputation: 96127

Here is one that works for arbitrary iterables:

>>> from itertools import islice
>>> [list(islice(it, n)) for it in [iter(lst)] for n in l]
[['a', 'a'], ['b', 'b', 'b'], ['c'], ['d', 'd']]
>>>

Although, probably would be cleaner as a generator:

def partition_by(iterable, partitions):
    it = iter(iterable)
    for n in partitions:
        yield list(islice(it, n))

And use like:

>>> list(partition_by(lst, l))
[['a', 'a'], ['b', 'b', 'b'], ['c'], ['d', 'd']]

Note, this approach will not error out and keep going if the iterator is exhausted:

>>> list(partition_by(lst, [10, 10, 10]))
[['a', 'a', 'b', 'b', 'b', 'c', 'd', 'd'], [], []]

If this behavior is undesirable:

def partition_by(iterable, partitions, strict=True):
    it = iter(iterable)
    for n in partitions:
        part = list(islice(it, n))
        if not part and strict:
            raise ValueError("iterable ran out of items early")
        yield part

Would do it.

Type annotated:

import typing
from collections.abc import Iterable

def partition_by(
    iterable: Iterable[T], 
    partitions: Iterable[int], 
    strict: bool=True
) -> Iterable[list[T]]:
    it = iter(iterable)
    for n in partitions:
        part = list(islice(it, n))
        if not part and strict:
            raise ValueError("iterable ran out of items early")
        yield part

Upvotes: -1

mozway
mozway

Reputation: 261580

For fun, here is a solution using the relatively new assignment expressions (python ≥ 3.8):

x=0
partitioned_lst = [lst[x:(x:=x+y)] for y in l]

NB. I wouldn't use this in production code as it depends on the outer scope

output:

[['a', 'a'], ['b', 'b', 'b'], ['c'], ['d', 'd']]

Upvotes: 0

user18098820
user18098820

Reputation:

lst=['a','a','b','b','b','c','d','d']
l=[2,3,1,2]
s = 0
Res=[]
for i in l:
  Res.append(lst[s: s+i])
  s += i
print(Res)

Output

[['a', 'a'], ['b', 'b', 'b'], ['c'], ['d', 'd']]

Upvotes: 1

Timus
Timus

Reputation: 11341

If you can't use groupby (see Peter Wood's comment) because you need l:

lst=['a','a','b','b','b','c','d','d']
l=[2,3,1,2]

it = iter(lst)
result = [[next(it) for _ in range(n)] for n in l]

Upvotes: 1

user3366507
user3366507

Reputation:

One liner just for fun

print([lst[sum(l[:index]):sum(l[:index])+num] for index, num in enumerate(l)])

Upvotes: 1

Salvatore Daniele Bianco
Salvatore Daniele Bianco

Reputation: 2691

partitioned_lst=[]
i=0
for n in l:
    partitioned_lst.append(lst[i:i+n])
    i+=n
partitioned_lst

Upvotes: 1

Related Questions