16

I know how to split a list into even groups, but I'm having trouble splitting it into uneven groups.

Essentially here is what I have: some list, let's call it mylist, that contains x elements.

I also have another file, lets call it second_list, that looks something like this:

{2, 4, 5, 9, etc.}

Now what I want to do is divide mylist into uneven groups by the spacing in second_list. So, I want my first group to be the first 2 elements of mylist, the second group to be the next 4 elements of mylist, the third group to be the next 5 elements of mylist, the fourth group to be the next 9 elements of `mylist, and so on.

Is there some easy way to do this? I tried doing something similar to if you want to split it into even groups:

for j in range(0, len(second_list)):
    for i in range(0, len(mylist), second_list[j]):
        chunk_mylist = mylist[i:i+second_list[j]]

However this doesn't split it like I want it to. I want to end up with my # of sublists being len(second_list), and also split correctly, and this gives a lot more than that (and also splits incorrectly).

Yu Hao
  • 119,891
  • 44
  • 235
  • 294
J. P.
  • 289
  • 4
  • 14

6 Answers6

26

You can create an iterator and itertools.islice:

mylist = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
seclist = [2,4,6]

from itertools import islice
it = iter(mylist)

sliced =[list(islice(it, 0, i)) for i in seclist]

Which would give you:

[[1, 2], [3, 4, 5, 6], [7, 8, 9, 10, 11, 12]]

Once i elements are consumed they are gone so we keep getting the next i elements.

Not sure what should happen with any remaining elements, if you want them added, you could add something like:

mylist = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13 ,14]
seclist = [2, 4, 6]

from itertools import islice

it = iter(mylist)

slices = [sli for sli in (list(islice(it, 0, i)) for i in seclist)]
remaining = list(it)
if remaining:
    slices.append(remaining)
print(slices)

Which would give you:

 [[1, 2], [3, 4, 5, 6], [7, 8, 9, 10, 11, 12], [13, 14]]

Or in contrast if there were not enough, you could use a couple of approaches to remove empty lists, one an inner generator expression:

from itertools import islice

it = iter(mylist)
slices = [sli for sli in (list(islice(it, 0, i)) for i in seclist) if sli]

Or combine with itertools.takewhile:

 from itertools import islice, takewhile

it = iter(mylist)
slices = list(takewhile(bool, (list(islice(it, 0, i)) for i in seclist)))

Which for:

mylist = [1, 2, 3, 4, 5, 6]
seclist = [2, 4, 6,8]

would give you:

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

As opposed to:

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

What you use completely depends on your possible inouts and how you would like to handle the various possibilities.

Ohad Eytan
  • 8,114
  • 1
  • 22
  • 31
Padraic Cunningham
  • 176,452
  • 29
  • 245
  • 321
  • 2
    Nice use of `islice` -- consuming different amounts of items from *the same* iterator. +1 – jedwards Aug 09 '16 at 22:54
  • wish i could upvote twice; such an extremely useful answer i put it into my personal general purpose tool library. – Rick Jun 13 '19 at 13:24
7

A numpythonic approach:

>>> lst = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
>>> sec = [2, 4, 5]
>>> np.split(lst, np.cumsum(sec))
[array([0, 1]), array([2, 3, 4, 5]), array([ 6,  7,  8,  9, 10]), array([11])]

And here is a Python3.X approach using itertool.accumulate():

>>> lst = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
>>> sec = [2,4,6]
>>> from itertools import accumulate
>>> sec = list(accumulate(sec_lst))
>>> sec = [0] + sec + [None] if sec[0] != 0 else sec + [None]
>>> 
>>> [lst[i:j] for i, j in zip(sec, sec[1:])]
[[0, 1], [2, 3, 4, 5], [6, 7, 8, 9, 10], [11]]
Mazdak
  • 105,000
  • 18
  • 159
  • 188
5

Using list-comprehensions together with slicing and sum() function (all basic and built-in tools of python):

mylist = [1,2,3,4,5,6,7,8,9,10]
seclist = [2,4,6]

[mylist[sum(seclist[:i]):sum(seclist[:i+1])] for i in range(len(seclist))]

#output:
[[1, 2], [3, 4, 5, 6], [7, 8, 9, 10]]

If seclist is very long and you wish to be more efficient use numpy.cumsum() first:

import numpy as np
cumlist = np.hstack((0, np.cumsum(seclist)))
[mylist[cumlist[i]:cumlist[i+1]] for i in range(len(cumlist)-1)]

and get the same results

Graham
  • 7,431
  • 18
  • 59
  • 84
Ohad Eytan
  • 8,114
  • 1
  • 22
  • 31
  • 1
    This was my initial solution too, but I'm thinking there must be a more efficient way than computing 2*len(seclist) sums. – jedwards Aug 09 '16 at 22:44
  • 1
    @jedwards one can build the [cumulative sum list](http://stackoverflow.com/questions/15889131/how-to-find-the-cumulative-sum-of-numbers-in-a-list) first if he worried about efficiency – Ohad Eytan Aug 09 '16 at 22:49
  • Yeah, that's where I was headed -- `seps = [sum(seclist[:i]) for i in range(len(seclist)+1)]; lists = [mylist[i:j] for (i,j) in zip(seps, seps[1:])]` – jedwards Aug 09 '16 at 22:51
1

This solution keeps track of how many items you've written. It will crash if the sum of the numbers in the second_list is longer than mylist

total = 0
listChunks = []
for j in range(len(second_list)):
    chunk_mylist = mylist[total:total+second_list[j]]
    listChunks.append(chunk_mylist)
    total += second_list[j]

After running this, listChunks is a list containing sublists with the lengths found in second_list.

nbryans
  • 1,507
  • 17
  • 24
  • @J.P., can you share what your `second_list` looks like? This code works fin in my test cases. – nbryans Aug 09 '16 at 22:52
  • 1
    Let me double check on my data (mylist is really long) and make sure - it made be a problem on my end in which case I think this would work! I'll double check on it and accept your answer once I do. Thanks! – J. P. Aug 09 '16 at 23:22
1
subgroups = []
start=0
for i in second_list:
  subgroups.append(mylist[start:start + i])
  start = i + start

At the end subgroups will contain the desired lists

Example run:

>>> mylist = [1,2,3,4,5,6,7,8,9,10,11,12]
>>> second_list = [2,4,5,9]
>>> subgroups = []
>>> start=0
>>> for i in second_list:
...    subgroups.append(mylist[start:start + i])
...    start = i + start
...
>>> subgroups
[[1, 2], [3, 4, 5, 6], [7, 8, 9, 10, 11], [12]]
Ankur Agarwal
  • 23,692
  • 41
  • 137
  • 208
1

Padraic has the best solution IMO, but I'll chime in with this hacky one liner that requires no imports:

>>> L = [1,2,3,4,5,6,7,8,9,10] # source list - could be any iterable
>>> S = [2,3,4] # group sizes - could be any iterable

>>> [ [ [ next(i) for k in range(j) ] for j in iter(S) ] for i in [iter(L)] ][0]
[[1, 2], [3, 4, 5], [6, 7, 8, 9]]
rovyko
  • 4,068
  • 5
  • 32
  • 44