split list comprehension into nested lists based on a rule

Question

I have a simple list splitting question:

given a nested list like this:

x = [[1,4,3],[2,3,5,1,3,52,3,5,2,1],[2]]

I want to further split any element(sub-list) longer than 3, and whose length is a multiple of 3 or 3n+1 into sub-lists of length 3, except for the last chunk, so the result I want is:

x2 = [[1,4,3], [2,3,5],[1,3,52],[3,5,2,1],[2]]

I think it can be done with itertools.groupby and/or yield functions... but couldn't put together the details >> a_function(x)...

 splits = [ a_function(x) if len(x)>3 and (len(x) % 3 == 0 or len(x) % 3 == 1) else x for x in x]

Could anyone kindly give me some pointers? Thanks so much.

For lists not divisible by three, how are they handled ... ? — g.d.d.c, Aug 01 '16 at 16:01
Also, http://stackoverflow.com/questions/312443/how-do-you-split-a-list-into-evenly-sized-chunks-in-python?rq=1. — g.d.d.c, Aug 01 '16 at 16:02
@g.d.d.c, thanks for your reply! 3n+2 values are left as they are... and thank you so much for the link!! tho I realized my question is a bit more involved as I re-edited just now... — shenglih, Aug 01 '16 at 16:06

Robᵩ · Accepted Answer · 2016-08-01T16:11:01.583

Sometimes, when the requirements for a list comprehension are unusual or complicated, I use a generator function, like so:

def gen(l):
   for sublist in l:
       if len(sublist)%3 == 0:
           for i in range(0, len(sublist), 3):
               yield sublist[i:i+3]
       elif len(sublist)%3 == 1:
           for i in range(0, len(sublist)-4, 3):
               yield sublist[i:i+3]
           yield sublist[-4:]
       else:
           yield sublist

# OP's data:
x = [[1,4,3],[2,3,5,1,3,52,3,5,2],[2]]
y = [[1,4,3],[2,3,5,1,3,52,3,5,2,1],[2]]

# Using either list comprehension or list constructor:
newx = [item for item in gen(x)]
newy = list(gen(y))

# Result:
assert newx == [[1, 4, 3], [2, 3, 5], [1, 3, 52], [3, 5, 2], [2]]
assert newy == [[1, 4, 3], [2, 3, 5], [1, 3, 52], [3, 5, 2, 1], [2]]

score 0 · Answer 2 · answered Aug 01 '16 at 17:01

# Create the list 
x = [[1,4,3],[2,3,5,1,3,52,3,5,2,1],[2]] 
test = []
# Loop through each index of the list
for i in range(len(x)):
#if the length of the specific index is 3 then go do this loop. 
    if len(x[i]) > 3:
        j = len(x[i])
        # This cuts through the loop every 3 steps and makes it a new list. 
        test = ([x[i][j:j+3] for j  in range(0, len(x[i]), 3)])
        # If the length of the last index is 1 then add it to the previous index of the new list. 
        if len(test[-1]) == 1: 
                test[-2] = test[-2] + test[-1]
                # pop deletes the last entry
                test.pop(-1)
                x[i] = test 
        else:
            x[i] = test

You then get the output:

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

split list comprehension into nested lists based on a rule

2 Answers2