MatplotlibとNetworkXを使った探索木の可視化

探索アルゴリズムで何かしらの解を求めようとする際に、その探索過程を可視化することは探索アルゴリズムの理解に大きく貢献してくれます。
そこでMatplotlibとNetworkXを使って探索木の可視化をする方法を実装したので、備忘録の意味も込めてそれの紹介しようと思います。

環境・ライブラリ

• python: 3.6
• matplotlib: 3.0.1
• networkx: 2.4
• pydot: 1.4.1
• pygraphviz: 1.5
• jupyter: 1.0.0

pythonの環境構築がまだの方は「Python環境構築(pyenv + pyenv-virtualenv)」参考にしてください。

探索木のサンプルデータ

「sample_search_tree.csv」をクリックしてダウンロードしてください。
1行に「あるノード,それに対応する親ノード」の識別子が記載してあります。
この二つの情報を保存するようにしておけば自分の探索木でも可視化することができます。
以下にファイルの中身の一部を記載しておきます。

sample_search_tree.csv

0,1
2,1
3,2
...
100,99

探索木の可視化

可視化結果

まずどんな感じで描画されるのか、先に実際の可視化した探索木を見せます。

サンプルコード

※Jupyter notebookにて実行しています。

立ち上げたJupyter notebookと同じ階層に先ほどダウンロードしたsample_search_tree.csvを配置するようにしてください。

visualize_search_tree.ipynb

import csv

import matplotlib.pyplot as plt
import networkx as nx
from networkx.drawing.nx_agraph import graphviz_layout

# read the sample file
with open("sample_search_tree.csv", 'r') as f:
    logs = [r for r in csv.reader(f)]

# networkx settings
G = nx.DiGraph()
for log in logs:  #  log=[self node, parent node]
    G.add_node(log[0])
    G.add_edge(log[1], log[0])
p = nx.drawing.nx_pydot.to_pydot(G)
pos = graphviz_layout(G, prog="dot")

# matplotlib settings
fig = plt.figure(figsize=(10, 10), dpi=300)
ax = fig.add_subplot(1, 1, 1)

# draw the network
nx.draw(G,
        ax=ax,
        pos=pos,
        with_labels=False,
        arrows=False,
        node_size=100,
        node_shape='o',
        width=0.5,
        node_color=range(len(logs)+1),
        cmap="jet")

# draw the color map bar
sm = plt.cm.ScalarMappable(cmap="jet_r", norm=plt.Normalize(vmin=0, vmax=1))
sm._A = []
cbar = fig.colorbar(sm, shrink=0.4, ticks=[1, 0])
cbar.ax.set_yticklabels(['start', 'end'], fontdict={'fontsize': 20})

ちょっとだけ解説

普通にNetworkを描画しようとすると木構造にはならないので、G = nx.DiGraph()(有向グラフ)をベースにしてgraphviz_layoutで体裁を整えるようにしています。
そしてnx.draw()での出力の際に、ノードラベル・エッジの矢印があると見にくいので非表示にしています。

探索木の可視化では、どういう順番に探索されたかは重要な情報なため探索順に色が変化するようにnode_colorとcmapを設定しています。
今回の図では、青色（探索初期）〜赤色（探索終盤）となるように色付けをしています。

終わりに

簡単ではありますが、探索木の可視化をする方法について備忘録的にまとめてみました。
探索アルゴリズムについて勉強する際には、スクリプトからは再帰関数の挙動がいまいちわからないと思います。
そこで実際にどのような動きをしているのかを可視化することで、どのようにアルゴリズムが動いてるか・その挙動が期待しているものなのかをざっくりと確認することができるようになります。
この可視化スクリプトが探索アルゴリズムを使っている人の少しでも助けになれば幸いです。

それでは。

参照

• networkx.drawing.nx_pylab.draw_networkx
• Node Colormap
• How to add a simple “colorbar” to a network-graph plot in python?