Pandas: Looking to create a multiple nested dictionary-CodePudding

Here is what I am looking to generate:

{A: {1: [1,2], 2: [2,5]},
 B: {3: [1,4], 4: [7,8]}}

Here is the df:

id sub_id
A  1
A  2
B  3
B  4

and I have the following array:

[[1,2],
[2,5],
[1,4],
[7,8]]

So far, I have the following code:

sub_id_array_dict = dict(zip(df['sub_id'].to_list(), arrays))

this results in the following dictionary:

{1: [1,2],
 2: [2,5],
 3: [1,4],
 4: [7,8]}

Now, I feel like i've gone down the wrong path as I'm not sure how to get roll it up to the id level.

Any help would be much appreciated.

CodePudding user response：

With a simple loop this can be done like:

from collections import defaultdict

sub_id_array_dict = defaultdict(dict)
for i, s, a in zip(df['id'].to_list(), df['sub_id'].to_list(), arrays):
    sub_id_array_dict[i][s] = a

CodePudding user response：

You can use pivot:

df['col'] = your_array
out = {k: {kk:vv for kk,vv in v.items() if vv is not np.nan} for k,v in df.pivot("sub_id","id",'col').to_dict().items()}

You can also do the same job using dict.setdefault:

out = {}
for i, j, l in zip(df['id'].to_list(), df['sub_id'].to_list(), your_array):
    out.setdefault(i, {}).setdefault(j, {})
    out[i][j] = l

Output:

{'A': {1: [1, 2], 2: [2, 5]}, 'B': {3: [1, 4], 4: [7, 8]}}