CODEWITHSUNDEEP
pythonpython-3.xgroupingpython-itertools
Igor Dragushhak
Igor Dragushhak

Reputation: 637

How to groupby a two dimentional list by a key with a condition?

I have this code:

from itertools import groupby
a = [[1,'n'],[2,'n'],[3,'n'],[4,'d'],[5,'n']]
b = [list(group) for key, group in groupby(a, lambda x: x[1] if x[1]=='n' else None)]
print(b)

Output:

[[[1, 'n'], [2, 'n'], [3, 'n']], [[4, 'd']], [[5, 'n']]]

Expected output:

[[[1, 'n'], [2, 'n'], [3, 'n']], [[5, 'n']]]

Upvotes: 1

Views: 73

Answers (2)

Patrick Artner
Patrick Artner

Reputation: 51643

This finds the longest group, without knowing 'n' to be it:

a = [[1,'n'],[2,'n'],[3,'n'],[4,'d'],[5,'n'],[6,'q'],[6,'q'],[6,'q'],[6,'q'],[6,'q']]

k = max( (list(grp) for key,grp in groupby(a,lambda x:x[1])), key=len)  
print(k)

Output:

[[6, 'q'], [6, 'q'], [6, 'q'], [6, 'q'], [6, 'q']]

If there are two groups of same length, the first one (irrespect. of 2nd element) is outputted.

This is a "unrolled" variant spezialized in finding the longest group of elements with 'n' as second element - using max():

a = [[1,'n'],[2,'n'],[3,'n'],[4,'d'],[5,'n']]

result = [[]]
for e in a:
    if e[1] == 'n':             # if 'n' 2nd element, add to list
        result[-1].append(e)
    elif result[-1]:
        result.append([])       # if not, append new sublist

longest = max(result, key=len)  # get the longest one from all lists

print(longest)

Output:

[[1, 'n'], [2, 'n'], [3, 'n']]

Doku:

Upvotes: 0

tobias_k
tobias_k

Reputation: 82899

You put the if in the wrong place. It should be a filter to the list comprehension, not part of the groupby key function.

>>> [list(group) for key, group in groupby(a, lambda x: x[1]) if key=='n']
[[[1, 'n'], [2, 'n'], [3, 'n']], [[5, 'n']]]

Upvotes: 5

Related Questions