从头开始构建决策树后,使用字典作为数据结构,该结构可以具有字典列表或唯一值。
但是,我在遍历它以使用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)