0%

如何对一棵树进行可视化(python anytree 简易实现)

Python 中有个很酷的第三方包叫做 anytree,全名 Any Python Tree Data,i.e. 期望用来表示任何树的数据结构。

而其中的可视化功能,每次都令人印象深刻。这篇文章简单分享,个人解决问题的思考路径 & 简易实现~

1
2
3
4
5
6
7
8
9
10
>> ...
>> r = RenderTree(root)
>> print(r)
Node('/A')
├── Node('/A/B')
│ ├── Node('/A/B/D')
│ │ └── Node('/A/B/D/F')
│ │ └── Node('/A/B/D/F/G')
│ └── Node('/A/B/E')
└── Node('/A/C')

1 问题拆解

一开始看到这个问题,可能有些没有头绪,但有没有可能对该问题进行分解🤔
.
.
.
.
.
.
.
.
.

一棵树的可视化 ,分解为:

  1. 每一行的显示,由三个部分组成
    1. 填充()
    2. 前缀(├── or └──)
    3. 节点自身
  2. 从上至下打印的顺序(深度优先遍历)

2 实现

2.1 定义「行」的数据结构

行(Row)与节点一一对应,其中包含两个元素:

  1. node 代表树的节点
  2. continues 中每个元素表示:根节点至当前节点路径中的每个节点,是不是对应层级中的最后一位。
用来生成每行的前缀
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from dataclasses import dataclass
from typing import List

from anytree import Node

LAST_NODE_PRE = "╰──"
NODE_PRE = "├──"

INDENT = "│"
BLANK = ""


@dataclass
class Row:
node: Node
continues: List[int]

@property
def pre(self):
# 根结点特殊处理
if len(self.continues) == 0:
return ""

indent = "".join([INDENT if x else BLANK for x in self.continues[:-1]])
branch = NODE_PRE if self.continues[-1] else LAST_NODE_PRE
return indent + branch

def __str__(self):
return self.pre + self.node.name

2.2 深度遍历(DFS)

一般会使用 stack 先入后出的特性,这里简单利用生成器,即 yield 关键字实现一版:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from typing import Tuple
from anytree import Node
from row import Row


def dfs(root: Node, continues: Tuple[bool] = None):
if continues is None:
continues = tuple()

yield Row(root, continues)

if not root.children:
return
for child in root.children:
is_last = root.children[-1] == child
for grandchild in dfs(child, continues + (not is_last,)):
yield grandchild

2.3 初始化树 & 打印

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from anytree import Node, RenderTree
from render import dfs


def init_tree():
a = Node("A")
b = Node("B", parent=a)
c = Node("C", parent=a)

d = Node("D", parent=b)
e = Node("E", parent=b)

f = Node("F", parent=d)
g = Node("G", parent=f)

return a


if __name__ == '__main__':
root = init_tree()
for x in dfs(root):
print(x)
"""
A
├── B
│ ├── D
│ │ ╰── F
│ │ ╰── G
│ ╰── E
╰── C
"""

总结

具体的实现在对 问题拆解 后,设计好对应的 数据结构和算法,就能比较容易清晰的实现:

当然你也可以直接查看 anytree 的源码:anytree.render.RenderTree,其中有更多可扩展性的设计,例如样式主题,打印的模式等。

🍻cheers! have a gooood day~~