defaultdict 嵌套用法

defaultdict 嵌套用法

起步

昨天在公司加班要处理这样一组数据：

{"村": "李村", "症状": "血糖", "姓名": "小李1"}
{"村": "李村", "症状": "血糖", "姓名": "小李2"}
{"村": "李村", "症状": "体量", "姓名": "小李3"}
{"村": "李村", "症状": "空腹", "姓名": "小李4"}
...

{"村": "刘村", "症状": "体量", "姓名": "小刘1"}
{"村": "刘村", "症状": "血糖", "姓名": "小刘2"}
{"村": "刘村", "症状": "空腹", "姓名": "小刘3"}
...

{"村": "王村", "症状": "空腹", "姓名": "小王1"}
{"村": "王村", "症状": "血糖", "姓名": "小王2"}
{"村": "王村", "症状": "体量", "姓名": "小王3"}
...

当时的需求是，统计出每个村每种症状的人数。数据是放在一个列表中一起返回的。

方法1：初始化目标数据结构

第一个想到的办法是手动初始化一个数据结构如下：

count = {
  "李村": {
    "血糖": 0,
    "体量": 0,
    ...
  },
  "刘村": {
    ...
  },
  ...
}

这样一来，遍历样例数据就可以直接做累加操作：

for item in sample_data:
    count[item["村"]][item["症状"]] += 1

可是我事先并不清楚有哪些村，有哪些病症，没办法初始化 count 这样的数据结构出来。当然也并非不可以，只要先遍历一次数据结构，用 set() 记录下村和症状的值即可。显然，在当时我不愿意这么做。

方法2：setdefault

dict.setdefault(k, new_v) 方法的作用是：如果存在键 k，就返回 k 对应的 v 值；如果没有 k，就给字典添加键 k，对应值为 new_v，同时返回 k 对应的 new_v。

其逻辑很简单，我们也可以实现：

def my_setdefault(dic, k, v):
    if k in dic:
        pass
    else:
        dic[k] = v
    return dic[k]        

利用 setdefault 要比用 if…else… 看起来简洁许多：

count = {}
    for item in sample_data:
        vill = item["村"]
        ill = item["症状"]

        count.setdefault(vill, {})
        count[vill].setdefault(ill, 0)

        count[vill][ill] += 1

只是代码看起来还是累赘。事实上昨天项目要得急，一时间也想不到更好的法子，也就用了这个方法。

方法3：collections.defaultdict

昨天项目上用 setdefault 之前，我考虑过 collections.defaultdict，因为用它的确很方便，不需要自己去初始化值，可以直接累加。但当时的需求是嵌套字典，于是我尝试 count = defaultdict(defaultdict(int))，结果当然是不可以。程序抛出异常：

TypeError: first argument must be callable or None

今天我就在想，这个错误明显是外层的 defaultdict 抛出来的，因为 defaultdict(int) 是不可以调用的。那我让它可调用不就行了吗。于是 partial 偏函数登场。

count = defaultdict(
    partial(defaultdict, int)
)
for item in sample_data:
    count[item["村"]][item["症状"]] += 1

我抱着试一试的心态点击 run，程序运行正常。

也就说，借助 partial 函数我们就可以 defaultdict 嵌套使用，处理上述相似需求时，代码能够简洁许多。