Praveen
Praveen

Reputation: 9345

Python convert list into nested list based on condition

I have a list

data = [1, 2, 3, 0, 0, 4, 5, 0, 0, 0, 4, 6, 7, 0]

expected output

result = [[1,2,3], [4,5], [4,6,7]]

Split the list when a zero occur; add all the non-zero numbers to a new list till the next zero

This is what I have tried:

res = []
tmp = []
for i in data:
  if i == 0 and len(tmp) > 0:
    res.append(tmp)
    tmp = []
  elif i != 0:
    tmp.append(i)

I am wondering is there a pythonic way of doing the same...

Upvotes: 2

Views: 126

Answers (2)

kederrac
kederrac

Reputation: 17322

you can use more_itertools.split_at:

from more_itertools import split_at

data = [1, 2, 3, 0, 0, 4, 5, 0, 0, 0, 4, 6, 7, 0]

list(filter(None, split_at(data, lambda x: x==0 )))

output:

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

Upvotes: 1

user2390182
user2390182

Reputation: 73470

You can use itertools.groupby:

from itertools import groupby

data = [1, 2, 3, 0, 0, 4, 5, 0, 0, 0, 4, 6, 7, 0]

result = [list(g) for k, g in groupby(data, key=bool) if k]
# [[1, 2, 3], [4, 5], [4, 6, 7]]

The grouping key function uses the fact that bool maps 0 to False and all other ints to True. You could be more explicit and use

key=lambda n: n != 0

Upvotes: 4

Related Questions