Jadox
Jadox

Reputation: 157

How to make a list comprehension in python with unequal sublists

I have a list of unequal lists. I would like to generate a new list with list comprehension from the sublists.

s = [['a','b','c','d'],['e','f','g'],['h','i'],['j','k','l','m']]

I am trying the following code but it keeps raising an indexError:

new_s = []
for i in range(len(s)):
    new_s.append((i,[t[i] for t in s if t[i]))

The expected output would be:

new_s = [(0,['a','e','h','j']),(1,['b','f','i','k']),(2,['c','g','l']),(3,['d','m'])]

Any ideas how to get this to work?

Upvotes: 10

Views: 449

Answers (3)

VPfB
VPfB

Reputation: 17267

This is without itertools, but also without a comprehension, so I don't know if it does count as a solution.

s = [['a','b','c','d'],['e','f','g'],['h','i'],'j','k','l','m']]

new_s = []
for i in range(len(s)):
    tmp = []
    for item in s:
        tmp.extend(item[i:i+1])
    new_s.append((i, tmp))

Upvotes: 0

Totoro
Totoro

Reputation: 887

one without itertools but modifying the original list.

def until_depleted():
   while any(sl for sl in s):
    yield 1

list(enumerate(list(s_l.pop(0) for s_l in s if s_l) for _ in until_depleted()))

one without modifying the original but with a counter

idx = 0
max_idx = max(len(_) for _ in s)

def until_maxidx():
    global idx
    while idx < max_idx:
        yield 1
        idx += 1

list(enumerate(list(s_l[idx] for s_l in s if idx < len(s_l)) for _ in until_maxidx()))

A more explicit one without the inner comprehension nor calling generators:

ret = []
idx = 0
max_idx = max(len(_) for _ in s)
while idx < max_idx:
    ret.append(list(s_l[idx] for s_l in s if idx < len(s_l)))
    idx += 1

print(list(enumerate(ret)))

Upvotes: 0

Cory Kramer
Cory Kramer

Reputation: 117866

You can use itertools.zip_longest to iterate over each sublist elementwise, while using None as the fill value for the shorter sublists.

Then use filter to remove the None values that were used from padding.

So all together in a list comprehension:

>>> from itertools import zip_longest
>>> [(i, list(filter(None, j))) for i, j in enumerate(zip_longest(*s))]
[(0, ['a', 'e', 'h', 'j']), (1, ['b', 'f', 'i', 'k']), (2, ['c', 'g', 'l']), (3, ['d', 'm'])]

Upvotes: 13

Related Questions