How to create sub-lists with 3 elements, but if a sub-list gets a single element then instead create two with 2 elements?

Viewed 237

Follows an example to clarify the 2-elements sub-lists use case:

If a list has 10 elements [A, B, C, D, E, F, G, H, I, J] I would like to get 4 sub-lists: [[A, B, C], [D, E, F], [G, H], [I, J]].

Thanks in advance.

3 Answers

I would do a function:

def partition(a):
    arr = np.array(a)
    n = len(arr) # assume that n>=3

    r = n % 3
    if r == 0:
        return list(arr.reshape(-1, 3))
    if r == 1:
        return list(arr[:-4].reshape(-1,3)) + list(arr[-4:].reshape(-1,2))

    # r == 2
    return list(arr[:-2].reshape(-1,3)) + [arr[-2:]]

Test:

partition(list('ABCDEFGHIJ'))
# [array(['A', 'B', 'C'], dtype='<U1'), array(['D', 'E', 'F'], dtype='<U1'),  
# array(['G', 'H'], dtype='<U1'), array(['I', 'J'], dtype='<U1')]

partition(list('ABCDEFGHI'))
# [array(['A', 'B', 'C'], dtype='<U1'), array(['D', 'E', 'F'], dtype='<U1'),
# array(['G', 'H', 'I'], dtype='<U1')]

partition(list('ABCDEFGH'))
# [array(['A', 'B', 'C'], dtype='<U1'), array(['D', 'E', 'F'], dtype='<U1'),
# array(['G', 'H'], dtype='<U1')]

How about:

l = list('ABCDEFGHIJ')
n = len(l)
res = [ l[(i if i < n-1 else i-1) : (i+3 if i < n-4 or i == n-3 else i+2)] for i in range(0,n,3)]
print(res)

This would give following outputs:

[['A', 'B', 'C'], ['D', 'E', 'F'], ['G', 'H', 'I']]
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G', 'H'], ['I', 'J']]
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G', 'H', 'I'], ['J', 'K']]
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G', 'H', 'I'], ['J', 'K', 'L']]
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G', 'H', 'I'], ['J', 'K'], ['L', 'M']]

Here's a solution using an iterator:

L = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"] 
N = len(L)
I = iter(L) 
[[next(I) for _ in range(size)] for size in [3]*(N//3-1) + ([3] if (N%3)==0 else ([2,2] if (N%3)==1 else [3,2]))] 
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G', 'H'], ['I', 'J']]
Related