こここブログ

機械学習, 統計分析, 競技プログラミング, CTF, VR, ……テクノロジーへの憧憬に身を焼かれている人のメモ帳です

Pythonで横向きのツリー図(樹形図)を描写

2019-07-19 公開 / 2019-07-20 更新

階層的なデータを可視化する際、ツリー図(樹形図)をPythonで描きたいこともある。

そんなとき、graphvizを使うと便利。

今回はデモとして東京都の区市町村名をツリー図で描写してみる。

入力と出力

入力

入力データは以下のようなデータ。

f:id:kokokocococo555:20190720165441p:plain

出力

出力はこうなる(一部掲載)。

f:id:kokokocococo555:20190720173440p:plain

元のcsvファイルをそのままグラフ化したような感じ。

環境構築

graphviz.readthedocs.io

スクリプト

PythonのgraphvizパッケージのAPIを使用してDOT言語のコードを生成し、そのコードをレンダリングするといった流れ。ノードやエッジが設定できる。

DOT言語等、詳しくは参考のサイトなどを参照。

横向きの図にするポイントはdot.attr(rankdir="LR")

まずは準備。

import pandas as pd
import numpy as np
from graphviz import Digraph

# # データをロード
csv_path = r"data/data1.csv"
df = pd.read_csv(csv_path, encoding="shift-jis")

次に、graphvizのAPIを使用してdot言語のスクリプトを作成していく。

各階層から順番にforループで情報を抽出していく。

# # graphvizのAPIを使用してdot言語のスクリプト作成------------
# ## インスタンスを作成
dot = Digraph(comment="tree")
dot.format = "svg"  # 保存するフォーマット
# ## フォント設定
# PNG出力の際に日本語を表示するために必要
dot.attr('node', fontname="Meiryo UI")
# ## グラフ設定
dot.attr(rankdir="LR")  # グラフを横方向に設定
dot.attr("node", shape="box", width="1", color="black")  # ノードのスタイル

# ## ノード、エッジを作成
# ### 第1階層------------
# 第1階層を一意に変更
root_set = df.iloc[:, 0].drop_duplicates()
for i in range(len(root_set)):
    # ルートを1つ抽出
    root = root_set.iloc[i]
    # ルートノードを作成
    dot.node("A{}".format(i), root)  # ルートの識別用にA1, A2などの名前を付ける
    # ルートの子となる行を抜き出し
    df_sub1 = df[df.iloc[:, 0]==root]
    # ### 第2階層------------
    # 第2階層を一意に変更
    node1_set = df_sub1.iloc[:, 1].drop_duplicates()
    for j in range(len(node1_set)):
        # 第2階層のノードを1つ抽出
        node1 = node1_set.iloc[j]
        # ノードを作成
        dot.node("B{}-{}".format(i, j), node1)  # ルートの識別用にB1-1, B1-2などの名前を付ける
        # ルートの子となる行を抜き出し
        df_sub2 = df_sub1[df_sub1.iloc[:, 1]==node1]
        # エッジを作成
        dot.edge("A{}".format(i), "B{}-{}".format(i, j))  # 第1階層と第2階層をつなぐ
        # ### 第3階層(最終階層)------------
        for k in range(len(df_sub2)):
            leaf = df_sub2.iloc[k, 2]
            # ノードを作成
            dot.node("C{}-{}-{}".format(i, j, k), leaf)

            # エッジを作成
            dot.edge("B{}-{}".format(i, j), "C{}-{}-{}".format(i, j, k))  # 第2階層と第3階層をつなぐ

# # 保存と表示
dot.render("output/tree.gv", view=True)

# (参考)
print(dot.source)  # dot言語での記述を表示

dot.attrを変更するなどしてノードやエッジの見た目を変えていい感じにもできる。

ソースコードはこちらにも。

github.com

東京都の人口をツリー図に加えて描写

入力データに人口の列が加わると、以下のように図に人口を表示することもできる(一部掲載)。

f:id:kokokocococo555:20190720182358p:plain

以下のようにPython上で数値処理を行い、ノード作成部分でノードのラベルに人口を加えている。

    # ルートを1つ抽出
    root = root_set.iloc[i]
    # ルートの子となる行を抜き出し
    df_sub1 = df[df.iloc[:, 0]==root]
    # 人口を計算
    population_of_root = df_sub1.iloc[:, -1].sum()
    # ルートノードを作成
    dot.node("A{}".format(i), "{}\n({:,}人)".format(root, population_of_root))

参考

graphviz.readthedocs.io

vdeep.net

swdrsker.hatenablog.com

qiita.com

kyle-in-jp.blogspot.com