defaultdict 嵌套用法

文章目錄

• defaultdict 嵌套用法
• • 起步
  • 方法1：初始化目标資料結構
  • 方法2：setdefault
  • 方法3：collections.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))

，結果當然是不可以。程式抛出異常：

今天我就在想，這個錯誤明顯是外層的 defaultdict 抛出來的，因為 defaultdict(int) 是不可以調用的。那我讓它可調用不就行了嗎。于是 partial 偏函數登場。

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

我抱着試一試的心态點選 run，程式運作正常。

也就說，借助 partial 函數我們就可以 defaultdict 嵌套使用，處理上述相似需求時，代碼能夠簡潔許多。