使用python和graphviz从自定义词典结构中绘制树

时间:2019-05-14 13:03:33

标签: python graphviz dot

从头开始构建决策树后,使用字典作为数据结构,该结构可以具有字典列表或唯一值。

但是,我在遍历它以使用graphviz点包进行绘制时遇到了困难。

这是我的尝试,虽然还不完整-但欢迎提出任何建议。

from graphviz import Digraph

tree =  {"cu_rate_m6 <= 0.000227778":
[1,
{"ratio_pd_l3_f3 <= 1.25":
[1,
{
"rev_a6 = 4":
[1,
{"avg_bill_amt_l3 <= 19724.333335": [
{
"avg_bill_amt_l3 <= 19669.166665":
[1, 0]},
1]}]}]}]}


graph = Digraph()
graph.attr('node', shape='box')


def draw_graph(dg, stree):

    # check if type is dict
    if isinstance(stree, dict):
        for k, v in stree.items():

            node_name = k
            # check if type is list
            if isinstance(v, list):
                if len(v) == 2:
                    yes_node = v[0]
                    no_node = v[1]
                    dg.edge(node_name, yes_node)
                    dg.edge(node_name, no_node)
                else:
                    yes_node = v[0]
                    dg.edge(node_name, yes_node)

draw = draw_graph(graph, tree)

0 个答案:

没有答案