lbedogni
lbedogni

Reputation: 7997

Create nested dictionaries of arbitrary length

I have the following problem. I have a nested dictionary like the following

A = {'A':{'STATES':['0','1','2']}, 'B':{'STATES':['10','20']}}

What I want to build eventually is a nested dictionary with all the possible combinations of STATES. So for a known number of keys in A is trivial as this

_dict = {}
for s in A['A']['STATES']:
    if s not in _dict:
        _dict[s] = {}
    for s1 in A['B']['STATES']:
        _dict[s][s1] = 0

This gives

{'0': {'10': 0, '20': 0}, '1': {'10': 0, '20': 0}, '2': {'10': 0, '20': 0}}

which is what I want. However I do not know the number of keys in A beforehand. What it would be an efficient solution to to the same with an arbitrary number of elements in A?

EDIT

For instance with three elements I would have

{'0': {'10': {'100':0}, '20': {'100':0}, '1': {'10': {'100':0}, '20': {'100':0}, '2': {'10': {'100':0}, '20': {'100':0}}

Upvotes: 0

Views: 381

Answers (2)

Ajax1234
Ajax1234

Reputation: 71451

You can use recursion:

A = {'A':{'STATES':['0','1','2']}, 'B':{'STATES':['10','20']}}
def combos(d):
   return 0 if not d else {i:combos(d[1:]) for i in d[0]}

print(combos([j['STATES'] for j in A.values()]))

Output:

{'0': {'10': 0, '20': 0}, '1': {'10': 0, '20': 0}, '2': {'10': 0, '20': 0}}

With more than two keys:

A = {'A':{'STATES':['0','1','2']}, 'B':{'STATES':['10','20']}, 'C':{'STATES':['100']}}
print(combos([j['STATES'] for j in A.values()]))

Output:

{'0': {'10': {'100': 0}, '20': {'100': 0}}, '1': {'10': {'100': 0}, '20': {'100': 0}}, '2': {'10': {'100': 0}, '20': {'100': 0}}}

Upvotes: 0

crissal
crissal

Reputation: 2647

This problem is a little complex, but it can be splitted up in three parts:

  1. Parse all the values, mapping a list to every valid key.
  2. Get the list of all the combinations in the order of dictionary insertion.
  3. Translate the list of tuples into a nested dictionary, looping over the values inside the tuple itself - because we don't know its length.
import itertools

A = {'A':{'STATES':['0','1','2']}, 'B':{'STATES':['10','20']}, 'C':{'STATES':['100']}}

# 1. get your dictionary A, but reduced, so that
# for every key you have a list of states if the key "STATES" exists
Ared = {k: A[k]["STATES"] for k in A if A[k].get("STATES")}
print(Ared)  # {'A': ['0', '1', '2'], 'B': ['10', '20'], 'C': ['100']}

# 2. get all the combinations
combs = list(itertools.product(*Ared.values()))
print(combs)  # [('0', '10', '100'), ('0', '20', '100'), ('1', '10', '100'), ('1', '20', '100'), ('2', '10', '100'), ('2', '20', '100')]

# 3. translate them into a nested dictionary
d = dict()
for comb in combs:
    old_dict = d
    for i, key in enumerate(comb):
        if i == len(comb) - 1:
            old_dict[key] = 0
        elif not old_dict.get(key):
            old_dict[key] = {}
        old_dict = old_dict[key]
        
print(d)  # {'0': {'10': {'100': 0}, '20': {'100': 0}}, '1': {'10': {'100': 0}, '20': {'100': 0}}, '2': {'10': {'100': 0}, '20': {'100': 0}}}

Upvotes: 2

Related Questions