pyqt 可视化二叉树

1. 思路

• QGraphicsEllipseItem 来绘制节点和节点的值
• painter drawLine 绘制节点之间的连线 线的起始点未结点的圆心位置
• 节点位置计算
  • 可通过二叉树前序遍历 子节点位置 = 父节点位置 + 偏移

可视化如下序列

[3,9,20,null,null,15,7]

效果:

2. 代码

2.1 结点类

class MyQGraphicsEllipseItem(QGraphicsEllipseItem):

    left_node = None
    right_node = None
    
    def __init__(self, parent_node, pos, node_val):
        super(MyQGraphicsEllipseItem, self).__init__(QRectF(pos[0], pos[1], pos[2], pos[3]))
        self.pos = pos
        self.radius = int((pos[2]) / 2 )
        print(pos, self.radius)
        self.parent_node = parent_node
        self.node_val = str(node_val)
        self.sel_brush = QBrush(QColor("#F62A66"))
        self.sel_pen = QPen(QColor("#FFD933"))

        self.un_sel_brush = QBrush(QColor("#374955"))
        self.un_sel_pen =  QPen(QColor("#FFD933"))
        self.seled = False
        self.switch_pen_brush()

    def paint(self, painter:QPainter, item, widgt):
        super().paint(painter, item, widgt)
        # 绘制节点 val
        painter.drawText(self.boundingRect(), Qt.AlignmentFlag.AlignCenter, self.node_val)
        if self.parent_node is not None:
            # 绘制节点之间连线
            painter.drawLine(self.pos[0] + self.radius, self.pos[1]+self.radius, self.parent_node.pos[0]+self.radius, self.parent_node.pos[1]+self.radius)
    
    def switch_pen_brush(self):
        if self.seled:
            self.setBrush(self.sel_brush)
            self.setPen(self.sel_pen)
        else:
            self.setBrush(self.un_sel_brush)
            self.setPen(self.un_sel_pen)


    def mouseReleaseEvent(self, event):
        self.seled = not self.seled
        self.switch_pen_brush()
        self.update()
        global_ctx.set_sel_node(self)

    def mousePressEvent(self, event):
        pass