CODEWITHSUNDEEP
pythonpython-3.xlistmergelist-comprehension
Brian
Brian

Reputation: 25

Python merging list of lists with varying length

I want to write a function that merges a list of lists. The lists can be of varying lengths.

selected_cells = [[[1],[2],[3]], [[4],[5],[6]], [[7],[8]]]

My expected output would be [[1,4,7], [2,5,8], [3,6]]

I've written the following function

def merge(lst1, lst2): 
  print('Merging', lst1, 'and', lst2)
  return [(sub + [lst2[i][-1]]) for i, sub in enumerate(lst1) if i < len(lst2)]

for idx, cell in enumerate(selected_cells):
  if idx < (len(selected_cells)-1):
    print(idx)
    selected_cells[0] = merge(selected_cells[0], selected_cells[idx+1])
    print('Output', selected_cells[0])

However this outputs [[1, 4, 7], [2, 5, 8]]

I feel like I'm missing something simple here and it's driving me nuts. Any help would be appreciated

Upvotes: 2

Views: 58

Answers (2)

jferard
jferard

Reputation: 8180

As written in comments, you can use zip_longest from itertools module.

>>> selected_cells = [[[1],[2],[3]], [[4],[5],[6]], [[7],[8]]]
>>> from itertools import zip_longest
>>> L = list(zip_longest(*selected_cells))
>>> L
[([1], [4], [7]), ([2], [5], [8]), ([3], [6], None)]

And then flatten the tuples and remove the None values:

>>> [[x[0] for x in t if x] for t in L]
[[1, 4, 7], [2, 5, 8], [3, 6]]

Another option is to use a fold (functools.reduce):

>>> selected_cells = [[[1],[2],[3]], [[4],[5],[6]], [[7],[8]]]
>>> import functools
>>> functools.reduce(lambda acc, x: [acc[i] + (x[i] if i < len(x) else []) for i in range(len(acc))], selected_cells)
[[1, 4, 7], [2, 5, 8], [3, 6]]

That's perhaps less intuitive, though.

Upvotes: 1

marcos
marcos

Reputation: 4510

You can use zip_longest as stated in the comments, because zip will stop at the shortest iterable, for example:

from itertools import zip_longest

selected_cells = [[[1],[2],[3]], [[4],[5],[6]], [[7],[8]]]

output = [
    [cell[0] for cell in row if cell is not None] 
    for row in zip_longest(*selected_cells) 
]
print(output)
>>> [[1, 4, 7], [2, 5, 8], [3, 6]]

Upvotes: 1

Related Questions