Использование python & graphviz для рисования дерева из пользовательской структуры словаря - PullRequest
Новая
       17

Использование python & graphviz для рисования дерева из пользовательской структуры словаря

0 голосов
/ 14 мая 2019

После построения дерева решений с нуля, используя словарь в качестве структуры данных, которая может иметь список словаря или уникальные значения.

Но мне трудно пройти его, чтобы нарисовать его с помощью пакета graphviz dot.

Вот моя попытка, хотя и неполная, но любые предложения приветствуются. 

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)
...