Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.

I want to split a list using another list which contains the lengths of each split.

E.g:

>>> print list(split_by_lengths(list('abcdefg'), [2,1]))
... [['a', 'b'], ['c'], ['d', 'e', 'f', 'g']]
>>> print list(split_by_lengths(list('abcdefg'), [2,2]))
... [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]    
>>> print list(split_by_lengths(list('abcdefg'), [2,2,6]))
... [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
>>> print list(split_by_lengths(list('abcdefg'), [1,10]))
... [['a'], ['b', 'c', 'd', 'e', 'f', 'g']]
>>> print list(split_by_lengths(list('abcdefg'), [2,2,6,5]))
... [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]

As you can notice, if the lengths list does not cover all the list I append the remaining elements as an additional sublist. Also, I want to avoid empty lists at the end in the cases that the lengths list produces more elements that are in the list to split.

I already have a function that works as I want:

def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(islice(iterable, n))

def split_by_lengths(list_, lens):
    li = iter(list_)
    for l in lens:
        elems = take(l,li)
        if not elems:
            break
        yield elems
    else:
        remaining = list(li)
        if remaining:
           yield remaining

But I wonder if there is a more pythonic way to write a function such that one.

Note: I grabbed take(n, iterable) from Itertools Recipes

Note: This question is a repost from a stackoverflow question

share|improve this question

2 Answers 2

up vote 5 down vote accepted

This is my first answer that I gave on stackoverflow:

from itertools import islice

def split_by_lengths(seq, num):
    it = iter(seq)
    for n in num:
        out = list(islice(it, n))
        if out:
            yield out
        else:
            return   #StopIteration 
    remain = list(it)
    if remain:
        yield remain

Here I am not using a for-else loop because we can end the generator by using a simple return statement. And IMO there's no need to define an extra take function just to slice an iterator.

Second answer:

This one is slightly different from the first one because this won't short-circuit as soon as one of the length exhausts the iterator. But it is more compact compared to my first answer.

def split_by_lengths(seq, num):
    it = iter(seq)
    out =  [x for x in (list(islice(it, n)) for n in num) if x]
    remain = list(it)
    return out if not remain else out + [remain]
share|improve this answer
    
Agreed that take(n, iterable) is not necessary. –  VGonPa Apr 4 at 10:39

You can do it by just slicing the sequence — there's really no advantage to having a helper function liketake(). I also added a keyword argument to make the returning of anything remaining optional.

def split_by_lengths(sequence, lengths, remainder=True):
    last, SequenceLength = 0, len(sequence)
    for length in lengths:
        if last >= SequenceLength: return  # avoid empty lists
        adjacent = last + length
        yield sequence[last:adjacent]
        last = adjacent
    if last < SequenceLength and remainder:
        yield sequence[last:]

print list(split_by_lengths(list('abcdefg'), [2, 1]))
# [['a', 'b'], ['c'], ['d', 'e', 'f', 'g']]
print list(split_by_lengths(list('abcdefg'), [2, 2]))
# [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
print list(split_by_lengths(list('abcdefg'), [2, 2, 6]))
# [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
print list(split_by_lengths(list('abcdefg'), [1, 10]))
# [['a'], ['b', 'c', 'd', 'e', 'f', 'g']]
print list(split_by_lengths(list('abcdefg'), [2, 2, 6, 5]))
# [['a', 'b'], ['c', 'd'], ['e', 'f', 'g']]
share|improve this answer

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.