从数据框创建嵌套JSON

Question

data={'category':['Medical','Medical','Research','Medical','Research'], 
       'countrycode':['US','CAN','US','CAN','US'],
       'stateCode':['AK','AB','MO','NT','OK'],
       'statecount':[600,100,200,760,90]} 
df=pd.DataFrame(data)

这是我的输入数据框。我想生成一个嵌套的JSON输出，格式如下：

{
  'Medical' : { 
                'US' : {'AK':600,'OK':90}
                'CAN': {'AB':160, 'NT':760}
               },
  'Research' : { 
                'US' : {'MO':200,'OK':90}
               }
 }

我一直在尝试使用以下代码，但似乎无法把国家代码放在第二层。

df.groupby('category').apply(lambda x: x.groupby('countrycode').apply(lambda y:y[['stateCode','statecount']]).to_json(orient='records'))

任何帮助或指导都非常感谢，我是Python新手。

Answer 1

这里有一个使用 groupby 的高效方法：

out = {}
for (k1,k2), g in df.groupby(['category', 'countrycode']):
    out.setdefault(k1, {})[k2] = g.set_index('stateCode')['statecount'].to_dict()

另外，还有一种效率稍低但可能更灵活的方法：

{k1: {k2: g2.set_index('stateCode')['statecount'].to_dict()
      for k2, g2 in g.groupby('countrycode')}
 for k1, g in df.groupby('category')}

输出结果：

{'Medical': {'CAN': {'AB': 100, 'NT': 760},
             'US': {'AK': 600}},
 'Research': {'US': {'MO': 200, 'OK': 90}}}