CODEWITHSUNDEEP
pythonlistpython-itertools
Zero
Zero

Reputation: 1899

How to group numbers between two occurences of a number using itertools.groupby

I have a list that looks like this -

nums = [0,0,0,0,1,1,2,3,4,5,6,0,0,0,0,1,2,3,4,5,6,0,0,0,0]

I want to get the numbers between the 0s in the list. For this, I used the code below -

groups = list(itertools.groupby(nums, lambda item:item != 0))
groups = list(filter(lambda item:item[0], groups))
list(map(lambda item:list(item[-1]), groups))

But I am getting empty lists as the output -

[[], []]

My desired output is -

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

How can I do this using itertools.groupby?

Upvotes: 1

Views: 64

Answers (1)

j1-lee
j1-lee

Reputation: 13939

Try:

from itertools import groupby

nums = [0,0,0,0,1,1,2,3,4,5,6,0,0,0,0,1,2,3,4,5,6,0,0,0,0]

output = [list(g) for k, g in groupby(nums, lambda x: x != 0) if k]
print(output) # [[1, 1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6]]

The doc puts

The returned group is itself an iterator that shares the underlying iterable with groupby(). Because the source is shared, when the groupby() object is advanced, the previous group is no longer visible.

Therefore my guest is that the first line

groups = list(itertools.groupby(nums, lambda item:item != 0))

is the culprit. As list iterates over groupby object, each group in it also gets consumed. You can instead use

groups = [(k, list(g)) for k, g in groupby(nums, lambda x: x != 0)]

to store the result.

To see this:

groups = groupby(nums, lambda x: x != 0)
k1, g1 = next(groups)
print(k1, list(g1)) # False [0, 0, 0, 0]

groups = groupby(nums, lambda x: x != 0)
k1, g1 = next(groups)
k2, g2 = next(groups)
print(k1, list(g1), k2, list(g2)) # False [] True [1, 1, 2, 3, 4, 5, 6]

Upvotes: 2

Related Questions