階層構造を持ったdictionaryの作り方。

そんなの五秒で実装できるのでは？

後輩A君から「Pythonで階層構造を持った辞書を作りたい」という質問がきた。当然、
「そんなの五秒でできるジャン」
と思ったわけである。しかし結局、質問に答えるまでに30分。解説するのに30分ほど時間を費やしてしまった。また、実装しろと言われたら忘れてしまいそうなので、メモとして残しておきたい。

そもそも、後輩君の質問は単に2重構造、3重構造の辞書を作りたいということではなく、下記のような入力を受け取って

example.csv

A1,B1,C1,3
A1,B1,C2,1
A1,B1,C3,5
A1,B2,C1,4
A1,B2,C2,3
A1,B2,C3,1
A1,B3,C1,3
A1,B3,C2,2
A1,B3,C3,5
A2,B1,C1,3
A2,B1,C2,5
A2,B1,C3,3
A2,B2,C1,2
A2,B2,C2,1
A2,B2,C3,3
A2,B3,C1,4
A2,B3,C2,4
A2,B3,C3,5

自動的に以下のような階層的なdictionaryをつくりたいということであった。

{'A1': {'B1': {'C1': 1,
               'C2': 1,
               'C3': 3},
        'B2': {'C1': 3,
               'C2': 3,
               'C3': 4},
        'B3': {'C1': 2,
               'C2': 5,
               'C3': 5}},
 'A2': {'B1': {'C1': 4,
               'C2': 4,
               'C3': 1},
        'B2': {'C1': 1,
               'C2': 3,
               'C3': 3},
        'B3': {'C1': 4,
               'C2': 2,
               'C3': 3}}}

五秒じゃできないじゃん

で、これに案外てこずってしまった。まず思いついたのが、defaultdictを使う方法である

import collections
hoge = collections.defaultdict(lambda : collections.defaultdict(lambda : collections.defaultdict(int))

みたいにlambda式を使ってやれば、3重にネストされたdictを作成することができる。しかし。この方法は微妙だ。そもそも、事前に階層の数を知ってる必要があるし、階層の数が、要素によって違う場合もあるだろうし、あまり一般性がない。さらに、lambda式でdefault値を定義したdefaultdictはpickle化するのも大変だ。
で考えたのが以下の方法。

import pprint

def make_tree_dict(inputs):
    tree_dict = {}
    for i, ainput in enumerate(inputs):
        pre_dict = tree_dict
        for j, key in enumerate(ainput):
            if j == len(ainput)-2:
                pre_dict[key] = ainput[-1]
                break
            elif key not in pre_dict:
                pre_dict[key] = {} 
            else:
                pass 
            pre_dict = pre_dict[key] 
    return tree_dict

if __name__ == "__main__":
    pp = pprint.PrettyPrinter(width=10,compact=True)
    inputs = []
    with open("example.csv") as f:
        for line in f:
            line = line.rstrip().split(",")
            inputs.append(line)
    hoge = make_tree_dict(inputs) 
    pp.pprint(hoge) 

実際に上記のプログラムを動かすと、先に示したような階層的なdictのoutputを得ることができる。
一度も、tree_dictに直接代入が行われないのに、tree_dictの中身が更新されていくという、なんとも不思議な感じがするプログラムだが、これが動くのである。解説まで載せようと思ったが、時間がないのでまた今度に、、、、

ついでに上記のスクリプトは、下記のように要素ごとに階層の数が異なるinputに対しても適用できる。

example2.csv

A1,B1,C1,1
A1,B1,C2,D1,3
A1,B1,C3,5
A1,B2,C1,D1,5
A1,B2,C2,2
A1,B2,C3,5
A1,B3,C1,2
A1,B3,C2,D1,4
A1,B3,C2,D2,10
A1,B3,C3,2
A2,B1,C1,4
A2,B1,C2,D1,5
A2,B1,C3,5
A2,B2,C1,D1,6
A2,B2,C2,3
A2,B2,C3,D1,8
A2,B3,C1,2
A2,B3,C2,5
A2,B3,C3,4

でこんな感じのdictが得られる

example2_output

{'A1': {'B1': {'C1': '1',
               'C2': {'D1': '3'},
               'C3': '5'},
        'B2': {'C1': {'D1': '5'},
               'C2': '2',
               'C3': '5'},
        'B3': {'C1': '2',
               'C2': {'D1': '4',
                      'D2': '10'},
               'C3': '2'}},
 'A2': {'B1': {'C1': '4',
               'C2': {'D1': '5'},
               'C3': '5'},
        'B2': {'C1': {'D1': '6'},
               'C2': '3',
               'C3': {'D1': '8'}},
        'B3': {'C1': '2',
               'C2': '5',
               'C3': '4'}}}

私はアホだった。

冷静に考えれば、defaultdictを再起的に定義すればいいだけだったわということに3年経って気づいてしまった。。。

from collections import defaultdict
recursive_defaultdict = lambda: defaultdict(recursive_defaultdict)
hoge = recursive_defaultdict()
hoge["A"]["T"]["G"]["C"] = 1
print(hoge["A"]["T"]["G"]["C"])

あー。ちょっとアホだったなぁ。。。。