How to use itertools.groupby when the key value is in the elements of the iterable?

Question

To illustrate, I start with a list of 2-tuples:

import itertools
import operator

raw = [(1, "one"),
       (2, "two"),
       (1, "one"),
       (3, "three"),
       (2, "two")]

for key, grp in itertools.groupby(raw, key=lambda item: item[0]):
    print key, list(grp).pop()[1]

yields:

1 one
2 two
1 one
3 three
2 two

In an attempt to investigate why:

for key, grp in itertools.groupby(raw, key=lambda item: item[0]):
    print key, list(grp)

# ---- OUTPUT ----
1 [(1, 'one')]
2 [(2, 'two')]
1 [(1, 'one')]
3 [(3, 'three')]
2 [(2, 'two')]

Even this will give me the same output:

for key, grp in itertools.groupby(raw, key=operator.itemgetter(0)):
    print key, list(grp)

I want to get something like:

1 one, one
2 two, two
3 three

I am thinking this is because the key is within the tuple inside the list, when in fact the tuple gets moved around as one. Is there a way to get to my desired output? Maybe groupby() isn't suited for this task?

Answer 1

groupby clusters consecutive elements of the iterable which have the same key. To produce the output you desire, you must first sort raw.

for key, grp in itertools.groupby(sorted(raw), key=operator.itemgetter(0)):
    print key, map(operator.itemgetter(1), grp)

# 1 ['one', 'one']
# 2 ['two', 'two']
# 3 ['three']

Answer 2

From the docs:

The operation of groupby() is similar to the uniq filter in Unix. It generates a break or new group every time the value of the key function changes (which is why it is usually necessary to have sorted the data using the same key function). That behavior differs from SQL’s GROUP BY which aggregates common elements regardless of their input order.

Since you are sorting the tuples lexicographically anyway, you can just call sorted:

for key, grp in itertools.groupby( sorted( raw ), key = operator.itemgetter( 0 ) ):
    print( key, list( map( operator.itemgetter( 1 ), list( grp ) ) ) )

Answer 3

I think a cleaner way to get your desired result is this.

>>> from collections import defaultdict
>>> d=defaultdict(list)
>>> for k,v in raw:
...  d[k].append(v)
... 
>>> for k,v in sorted(d.items()):
...  print k, v
... 
1 ['one', 'one']
2 ['two', 'two']
3 ['three']

building d is O(n), and now sorted() is just over the unique keys instead of the entire dataset