使用Qt4绘制二叉平衡树

最近在研究树结构，突发奇想将树结构使用Qt绘制出来，二话不说，开始动手。

主要思路：文件avl_tree.h包含树结构的实现，文件paint_tree.cpp和paint_tree.h实现一个由QWidget继承的对象，作为画板，将一棵avl树绘制出来。

注：只求实现功能，不考虑代码美观及效率等问题。

代码在bitbucket上，有兴趣的可以看看。

https://bitbucket.org/realid/painttree

avl_tree.h

#ifndef AVLTREE_H
#define AVLTREE_H
#include <QPoint>
#include <QPainter>
#include <cstdlib>
#include <QString>

template<class T>
class AvlNode {
    public:
        AvlNode():Left(NULL),Right(NULL),Height(0){}
    public:
        T Element;
        AvlNode *Left;
        AvlNode *Right;
        int Height;
};

template<class T>
class AvlTree {
    public:
        AvlTree():pRoot(NULL){}void Insert(T X);
        AvlNode<T>* Delete(T X);
        void Draw(AvlNode<T>* pNode, QPoint pt, QPainter* painter, 
                int scale, int height, int radi);private:
        int Height(AvlNode<T>* Node);
        int Max(int cmpa, int cmpb);
        AvlNode<T>* InsertPri(AvlNode<T>* Node, T X);
        AvlNode<T>* SingleRotateWithLeft(AvlNode<T>* pNode);
        AvlNode<T>* SingleRotateWithRight(AvlNode<T>* pNode);
        AvlNode<T>* DoubleRotateWithLeft(AvlNode<T>* pNode);
        AvlNode<T>* DoubleRotateWithRight(AvlNode<T>* pNode);
    public:
        AvlNode<T>* pRoot;
};

template<class T>
int AvlTree<T>::Height(AvlNode<T>* pNode)
{
    if (pNode == NULL)
        return -1;
    else
        return pNode->Height;
}

template<class T>
int AvlTree<T>::Max(int cmpa, int cmpb)
{
    return cmpa > cmpb ? cmpa : cmpb;
}

template<class T>
AvlNode<T>* AvlTree<T>::SingleRotateWithLeft(AvlNode<T> *pNode)
{
    AvlNode<T> *pLeft = pNode->Left;
    pNode->Left = pLeft->Right;
    pLeft->Right = pNode;

    pNode->Height = Max(Height(pNode->Left), Height(pNode->Right)) + 1;
    pLeft->Height = Max(Height(pLeft->Left), pNode->Height) + 1;

    return pLeft;
}

template<class T>
AvlNode<T>* AvlTree<T>::SingleRotateWithRight(AvlNode<T> *pNode)
{
    AvlNode<T>* pRight = pNode->Right;
    pNode->Right = pRight->Left;
    pRight->Left = pNode;

    pNode->Height = Max(Height(pNode->Left), Height(pNode->Right)) + 1;
    pRight->Height = Max(Height(pRight->Right), pNode->Height) + 1;

    return pRight;
}

template<class T>
AvlNode<T>* AvlTree<T>::DoubleRotateWithLeft(AvlNode<T>* pNode)
{
    pNode->Left = SingleRotateWithRight(pNode->Left);
    return SingleRotateWithLeft(pNode);
}

template<class T>
AvlNode<T>* AvlTree<T>::DoubleRotateWithRight(AvlNode<T>* pNode)
{
    pNode->Right = SingleRotateWithLeft(pNode->Right);
    return SingleRotateWithRight(pNode);
}

template<class T>
AvlNode<T>* AvlTree<T>::InsertPri(AvlNode<T>* pNode, T X)
{
    if (pNode == NULL) {
        pNode = new AvlNode<T>;
        pNode->Element = X;
    } else if (X < pNode->Element) {
        pNode->Left = InsertPri(pNode->Left, X);
        if (Height(pNode->Left) - Height(pNode->Right) > 1) {
            if (X < pNode->Left->Element)
                pNode = SingleRotateWithLeft(pNode);
            else
                pNode = DoubleRotateWithLeft(pNode);
        }
    } else if (X > pNode->Element) {
        pNode->Right = InsertPri(pNode->Right, X);
        if (Height(pNode->Right) - Height(pNode->Left) > 1) {
            if (X > pNode->Right->Element)
                pNode = SingleRotateWithRight(pNode);
            else
                pNode = DoubleRotateWithRight(pNode);
        }
    }
    pNode->Height = Max(Height(pNode->Left), Height(pNode->Right)) + 1;
    return pNode;
}

template<class T>
void AvlTree<T>::Insert(T X)
{ 
    pRoot = InsertPri(pRoot, X);
}

template<class T>
void AvlTree<T>::Draw(AvlNode<T>* pNode, QPoint pt, 
        QPainter* painter, int scale, int height, int radi)
{
    QPoint end;
    QRect rect(QPoint(pt.x()-radi, pt.y()-radi), QPoint(pt.x()+radi, pt.y()+radi));
    QBrush brush(Qt::green, Qt::SolidPattern);
    QString str = QString("%1").arg(pNode->Element, 0, 10);
    QFont font;
    font.setPixelSize(radi-2*str.length());
    if (pNode != NULL) {
        if (pNode->Left != NULL) {
            end.setX(pt.x() - scale);
            end.setY(pt.y() + height);
            painter->setPen(Qt::black);
            painter->drawLine(pt, end);
            Draw(pNode->Left, end, painter, scale/2, height, radi); 
        }
        if (pNode->Right != NULL) {
            end.setX(pt.x() + scale);
            end.setY(pt.y() + height);
            painter->setPen(Qt::black);
            painter->drawLine(pt, end);
            Draw(pNode->Right, end, painter, scale/2, height, radi); 
        }
        painter->setBrush(brush);
        painter->setPen(Qt::black);
        painter->drawEllipse(pt, radi, radi);
        painter->setPen(Qt::red);
        painter->setFont(font);
        painter->drawText(rect, Qt::AlignCenter,  str);
    }
}
#endif

 paint_tree.h

#ifndef PAINTTREE_H
#define PAINTTREE_H

#include <QWidget>
#include "avl_tree.h"

class QPaintEvent;
class QPainter;

class PaintTree : public QWidget
{
    public:
        PaintTree();
        void paintEvent(QPaintEvent *event);
    public:
        AvlTree<int>* pTree;
};
#endif

paint_tree.cpp

#include <QPainter>
#include <QPoint>
#include "paint_tree.h"

PaintTree::PaintTree()
{
    resize(800, 600);
    setWindowTitle(tr("Paint Tree"));
}

void PaintTree::paintEvent(QPaintEvent *)
{
    QSize s = size();
    QPainter painter(this);
    painter.setRenderHint(QPainter::Antialiasing, true);
    pTree->Draw(pTree->pRoot, QPoint(s.width()/2, 25), 
            &painter, s.width()/4, (s.height()-50)/(pTree->pRoot->Height),
            20);
}

最终效果：

再来一张从0-10的