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 嵌套使用,處理上述相似需求時,代碼能夠簡潔許多。