Reputation: 7997
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
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
Reputation: 2647
This problem is a little complex, but it can be splitted up in three parts:
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