Python 动画——一个简单的遍历有向无环图的算法

前言

之前在玩三蓝一棕（3b1b）大佬用 python 写的数学动画库 manim 时，一直想写一个例子来实践一下这个库，但是苦于没有好的选题，也不会大佬的那些奇妙的数学问题，这个例子的编写一直处于 pending 状态

在不久前写有向无环图（DAG）相关的代码时，偶然想到一个问题：如何遍历一个 DAG 的所有路径。当时只是为了解决简单的布局问题，后来将这个问题与上述的动画例子联系起来，发现它作为 manim 的例子再合适不过了。于是在这里记录一下

简单算法

算法本身其实只是一个简单的深度遍历。其所用到的三个栈分别为：当前栈(current)、输出栈(output)、分支栈(fork)。当前栈记录当前遍历的节点，输出栈记录当前遍历的路径

将根节点作为当前遍历元素
遍历当前遍历元素并将其压入输出栈
把当前遍历元素的子节点压入当前栈：
- 若没有子节点，则执行步骤5
- 若有大于一个子节点，则每多压入一个子节点，就在分支栈中压入一个当前遍历元素，直到所有子节点都压入
弹出当前栈顶元素，进行遍历，将其作为当前遍历元素，重复步骤2
检查输出栈顶元素：
- 若等于分支栈顶元素，则均弹出，当前遍历的路径为所有弹出过的元素加上输出栈剩余的元素的逆序，输出并清空当前遍历路径
- 若不等于分支栈顶元素，则弹出并继续执行此步骤直到输出栈为空，遍历结束

安装 manim

安装 manim 的方法此处略过

编写动画组件

由于 manim 是一个纯粹的数学动画库，因此并不预先包含算法里的一些数据结构的组件。为了将上述算法可视化，首先需要实现编写一些可用的组件

图的节点扩展自圆组件，唯一标识为其标签

# objects/node.py
from manimlib import Circle, RED_E, Text


class Node(Circle):
    CONFIG = {
        "radius": .3,
        "color": RED_E,
        "fill_opacity": 1.0,
    }

    def __init__(self, text="", *vmobjects, **kwargs):
        super().__init__(*vmobjects, **kwargs)
        text_obj = Text(text, font_size=28).shift(self.get_center())
        self.add(text_obj)

    def get_text(self):
        return self.submobjects[-1].text

图的边直接使用箭头组件；

栈的边用线段画成

# objects/stack.py
from manimlib import VGroup, ORIGIN, RIGHT, UP, BOTTOM, Line, Text, Scene, ReplacementTransform, FadeIn, FadeOut
from typing import List
from objects.stack_obj import *


class Stack(VGroup):
    def __init__(self, scene: Scene,  position=ORIGIN, text="", **kwargs):
        self.position = position
        self.bottomLen = RIGHT * 0.7
        self.wallHeight = UP * 2
        self.objHeight = 0.3
        self.scene = scene
        bottom_end = position + self.bottomLen
        left_side = Line(position, self.wallHeight + position)
        bottom = Line(position, self.bottomLen + position)
        right_side = Line(bottom_end, bottom_end + self.wallHeight)
        text = self.pin_text(text)
        super().__init__(left_side, bottom, right_side, text, **kwargs)
        self.stack: List[StackObj] =  []

    def pin_text(self, text):
        text = Text(text, font_size=28)
        text.next_to(self.position + RIGHT * 0.35, BOTTOM * 0.1)
        return text

    def update_text(self, text):
        text = self.pin_text(text)
        self.scene.play(ReplacementTransform(self.submobjects[-1], text))
        self.remove(self.submobjects[-1])
        self.add(text)
        return text

    def do_shift(self, position):
        animates = [item.animate.shift(position) for item in self.stack]
        self.scene.play(self.animate.shift(position), *animates)

    def get_length(self):
        return len(self.stack)

    def get_stack(self):
        return self.stack

栈中元素扩展自矩形，唯一标识为其标签

# objects/stack_obj.py
from manimlib import Rectangle, RED_E, Text


class StackObj(Rectangle):
    CONFIG = {
        "fill_opacity": 1,
        "color": RED_E
    }

    def __init__(self, width, height, text="", *vmobjects, **kwargs):
        super().__init__(width, height, *vmobjects, **kwargs)
        text_obj = Text(text, font_size=20).shift(self.get_center())
        self.add(text_obj)

    def get_text(self):
        return self.submobjects[0].text

另外编写节点集、边集、图组件，将所有的节点和边组合起来，方便管理

考虑到压栈和弹栈的动画，栈组件还需要编写相应的动画方法

# objects/stack.py
def push(self, obj_name, *vmobjects, **kwargs):
    content = StackObj(self.bottomLen[0], self.objHeight, obj_name, *vmobjects, **kwargs).shift(
        self.position + self.wallHeight - [.0, self.objHeight / 2, .0] + self.bottomLen / 2)
    self.scene.play(FadeIn(content, run_time=0.4))
    self.scene.wait(0.5)
    self.scene.play(content.animate.shift(-self.wallHeight +
               [.0, self.objHeight * (len(self.stack) + 1), .0]))
    self.stack.append(content)

def pop(self):
    content = self.stack.pop()
    self.scene.play(content.animate.shift(self.wallHeight -
               [.0, self.objHeight * (len(self.stack) + 1), .0]))
    self.scene.play(FadeOut(content, run_time=0.4))
    return content

实现算法

首先用 python 实现上述算法

def get_last_index(name, node_list):
    for i in range(len(node_list)):
        if name == node_list[i]:
            return i
    return

def get_children(index, side_list, node_list):
    result = []
    for i in range(len(side_list)):
        if side_list[i][0] == index:
            result.append(node_list[side_list[i][1]])
    return result

def traverse(nodes, sides):
    forks = []
    output = []
    current = []

    if not len(nodes):
        return
    output.append(nodes[0])

    has_no_children = False
    while len(output):
        next_obj = output[-1]
        next_index = get_last_index(next_obj, nodes)
        if has_no_children:
            has_no_children = False
            print([int(item) for item in output])
            while len(output) and (
                len(forks) <= 0
                or output[-1]
                != forks[-1]
            ):
                transfer = output.pop()
            if len(forks):
                forks.pop()
        else:
            children = get_children(next_index, sides, nodes)
            if len(children):
                for i in range(len(children)):
                    current.append(children[i])
                    if i:
                        forks.append(next_obj)
            else:
                has_no_children = True
                continue
        if len(current):
            transfer = current.pop()
            output.append(transfer)

测试一下

nodes = [str(item + 1) for item in range(0, 8)]
sides = [[0, 2],
            [0, 1],
            [1, 4],
            [1, 3],
            [2, 5],
            [3, 7],
            [4, 7],
            [4, 6],
            [4, 5],
            [4, 3],
            [5, 6],
            [6, 7]]

traverse(nodes, sides)