Search code examples
pythonlist-comprehension

Can I group / aggregate elements in a list (or dict) comprehension?


I have a list of tuples, as shown below:

[
    (1, "red")
    (1, "red,green")
    (1, "green,blue")
    (2, "green")
    (2, "yellow,blue")
]

I am trying to roll up the data, so that I can get the following dict output:

{
    1: ["red", "green", "blue"]
    2: ["green", "yellow", "blue"]
}

Notes are: the strings of colours are combined for the primary key (the number), and then split into a list, and de-duped (e.g. using set).

I'd also like to do the inverse, and group by the colours:

{
    "red": [1],
    "green": [1, 2]
    "yellow": [2]
    "blue": [1, 2]
}

I can clearly do this by looping through all of the tuples, but I'd like to try and do it with list / dict comprehensions if possible.


Solution

  • You can use collections.defaultdict:

    >>> from collections import defaultdict
    >>> lis = [                            
        (1, "red"),
        (1, "red,green"),
        (1, "green,blue"),
        (2, "green"),
        (2, "yellow,blue"),
    ]
    >>> dic = defaultdict(set)       #sets only contain unique items
    for k, v in lis:
        dic[k].update(v.split(','))
    
    >>> dic
    defaultdict(<type 'set'>,
    {1: set(['blue', 'green', 'red']),
     2: set(['blue', 'green', 'yellow'])})
    

    Now iterate over dic:

    >>> dic2 = defaultdict(list)
    for k,v in dic.iteritems():
        for val in v:
            dic2[val].append(k)
    ...         
    >>> dic2
    defaultdict(<type 'list'>,
    {'blue': [1, 2],
     'green': [1, 2],
     'yellow': [2],
     'red': [1]})