How to zip three lists into a nested dict

Question

I have three lists as:

z1 = ['A', 'A', 'B', 'B']
z2 = ['k1', 'k2', 'k1', 'k2']
z3 = ['v1', 'v2', 'v3', 'v4']

and when I write:

print(dict(zip(z2, z3)))

this is my output:

{'k2': 'v4', 'k1': 'v3'}

And I expect this:

{'A':{'k1': 'v1', 'k2': 'v2'} , 'B':{'k1': 'v3', 'k2': 'v4'}}

Could you please tell me how can I get my expected result?

Answer 1

Here's a one-liner using itertools.groupby, but aside from being a single expression, it doesn't really provide any benefit over the default-dict solution provided by RoadRunner.

>>> from itertools import groupby
>>> from operator import itemgetter
>>> keyf = itemgetter(0)
>>> dict((k, dict(v2 for _,v2 in v)) for k, v in groupby(zip(z1, zip(z2,z3)), key=keyf))
{'A': {'k2': 'v2', 'k1': 'v1'}, 'B': {'k2': 'v4', 'k1': 'v3'}}

This is only as short as it is because it takes advantage of the fact that z1 is already sorted. If it isn't, you'll need to sort the output of zip using the same key function before passing it to groupby.

dict((k, dict(v2 for _,v2 in v))
       for k, v in groupby(sorted(zip(z1, zip(z2,z3)),
                                  key=keyf),
                           key=keyf))

Breaking down how it works...

1. zip(z1, zip(z2, ze)) creates the key-value pairs for the outer dict:

   [('A', ('k1', 'v1')),
    ('A', ('k2', 'v2')),
    ('B', ('k1', 'v3')),
    ('B', ('k2', 'v4'))]
   

2. groupby effectively pairs each key (A or B) with its tuples:

   [('A', <itertools._grouper object at 0x100f656d0>),
    ('B', <itertools._grouper object at 0x100f655d0>)]
   

   Each _grouper is an iterable containing all the key/value pairs with the same key.

3. dict(v2 for _,v2 in v) extracts just the key/value pairs from the _groupers, leaving behind the key, which we can already get from the first element of the tuples returned by groupby.

Answer 2

The function zip() can accept more than two iterables. So you can use zip(z1, z2, z3) instead of zip(z2, z3). However, you still need to group the items since simply wrapping dict() will not work as it can't handle nested dictionaries needed for the 3-tuples.

To group the items correctly, I would use a collections.defaultdict():

from collections import defaultdict

z1 = ['A', 'A', 'B', 'B']
z2 = ['k1', 'k2', 'k1', 'k2']
z3 = ['v1', 'v2', 'v3', 'v4']

d = defaultdict(dict)
for x, y, z in zip(z1, z2, z3):
    d[x][y] = z

print(d)
# defaultdict(<class 'dict'>, {'A': {'k1': 'v1', 'k2': 'v2'}, 'B': {'k1': 'v3', 'k2': 'v4'}})

The above works because defaultdict(dict) initializes a dictionary for non-existent keys. It handles the dictionary creation for keys for you.

Additionally, If you wrap the end result with dict:

print(dict(d))
# {'A': {'k1': 'v1', 'k2': 'v2'}, 'B': {'k1': 'v3', 'k2': 'v4'}}

Note: defaultdict is just a subclass of dict, so you can treat it the same as a normal dictionary.

Answer 3

For the sake of completeness, you can use dict.setdefault, avoiding the import at the cost of a tiny overhead of creating and returning an empty dictionary at each iteration.

d = {}
for x, y, z in zip(z1, z2, z3):
    d.setdefault(x,{})[y] = z

print(d)
# {'A': {'k1': 'v1', 'k2': 'v2'}, 'B': {'k1': 'v3', 'k2': 'v4'}}

---

Another solution (not recommended) is using itertools.groupby:

d = {}
for k, g in groupby(enumerate(zip(z2, z3)), key=lambda x: z1[x[0]]):
    _, b = zip(*g)
    d[k] = dict(b)

print(d)
# {'A': {'k1': 'v1', 'k2': 'v2'}, 'B': {'k1': 'v3', 'k2': 'v4'}}