AVL平衡二叉树 - 插入算法

static int rotate_LL(struct bitree_t *tree, struct bitree_node_t *A)
{
    struct bitree_node_t *B = A->lchild;

    A->lchild = B->rchild;
    B->rchild = A;

    A->balancefactor = 0;
    B->balancefactor = 0;

    B->parent = A->parent;
    A->parent = B;
    if (B->parent)
    {
        if (B->parent->lchild == A)
        {
            B->parent->lchild = B;
        }
        else
        {
            B->parent->rchild = B;
        }
    }
    else
    {
        tree->root = B;
    }

    if (A->lchild)
    {
        A->lchild->parent = A;
    }

    return 0;
}

static int rotate_RR(struct bitree_t *tree, struct bitree_node_t *A)
{
    struct bitree_node_t *C = A->rchild;

    A->rchild = C->lchild;
    C->lchild = A;

    A->balancefactor = 0;
    C->balancefactor = 0;

    C->parent = A->parent;
    A->parent = C;
    if (C->parent)
    {
        if (C->parent->lchild == A)
        {
            C->parent->lchild = C;
        }
        else
        {
            C->parent->rchild = C;
        }
    }
    else
    {
        tree->root = C;
    }

    if (A->rchild)
    {
        A->rchild->parent = A;
    }

    return 0;
}

static int rotate_LR(struct bitree_t *tree, struct bitree_node_t *A)
{
    struct bitree_node_t *B = A->lchild;
    struct bitree_node_t *Ex = B->rchild;

    A->lchild = Ex->rchild;
    B->rchild = Ex->lchild;
    Ex->lchild = B;
    Ex->rchild = A;

    /* 链接父结点 */
    Ex->parent = A->parent;
    A->parent = Ex;
    B->parent = Ex;
    if (Ex->parent)
    {
        if (Ex->parent->lchild == A)
        {
            Ex->parent->lchild = Ex;
        }
        else
        {
            Ex->parent->rchild = Ex;
        }
    }
    else
    {
        tree->root = Ex;
    }

    if (A->lchild)
    {
        A->lchild->parent = A;
    }

    if (B->rchild)
    {
        B->rchild->parent = B;
    }

    /* 重新计算平衡因子 */
    if (Ex->balancefactor == 1)
    {
        A->balancefactor = -1;
    }
    else
    {
        A->balancefactor = 0;
    }

    if (Ex->balancefactor >= 0)
    {
        B->balancefactor = 0;
    }
    else
    {
        B->balancefactor = 1;
    }

    Ex->balancefactor = 0;
    return 0;
}

static int rotate_RL(struct bitree_t *tree, struct bitree_node_t *A)
{
    struct bitree_node_t *C = A->rchild;
    struct bitree_node_t *Dx = C->lchild;

    A->rchild = Dx->lchild;
    C->lchild = Dx->rchild;
    Dx->lchild = A;
    Dx->rchild = C;

    /* 链接父结点 */
    Dx->parent = A->parent;
    A->parent = Dx;
    C->parent = Dx;
    if (Dx->parent)
    {
        if (Dx->parent->lchild == A)
        {
            Dx->parent->lchild = Dx;
        }
        else
        {
            Dx->parent->rchild = Dx;
        }
    }
    else
    {
        tree->root = Dx;
    }

    if (A->rchild)
    {
        A->rchild->parent = A;
    }

    if (C->lchild)
    {
        C->lchild->parent = C;
    }

    /* 重新计算平衡因子 */
    if (Dx->balancefactor == -1)
    {
        A->balancefactor = 1;
    }
    else
    {
        A->balancefactor = 0;
    }

    if (Dx->balancefactor <= 0)
    {
        C->balancefactor = 0;
    }
    else
    {
        C->balancefactor = -1;
    }

    Dx->balancefactor = 0;
    return 0;
}


static int insert_avl(struct bitree_t *tree, struct bitree_node_t *newnode)
{
    /* 利用栈实现 */
    struct bitree_node_t *stack[BITREE_MAX_HEIGHT] = {0};
    int top = -1;
    struct bitree_node_t *node = NULL;
    struct bitree_node_t *ptr = NULL;

    /* 找到待插入的位置 */
    node = tree->root;
    while (node)
    {
        if (tree->cmp(newnode, node) == 0)
        {
            return -1;
        }

        /* 记录所走路径 */
        ptr = node;
        stack[++top] = ptr;

        if (tree->cmp(newnode, node) < 0)
        {
            node = node->lchild;
        }
        else
        {
            node = node->rchild;
        }
    }

    /* 插入节点 */
    if (tree->cmp(newnode, ptr) < 0)
    {
        ptr->lchild = newnode;
    }
    else
    {
        ptr->rchild = newnode;
    }
    newnode->parent = ptr;

    /* 遍历栈，使树平衡 */
    node = newnode;
    while (top >= 0)
    {
        ptr = stack[top--];

        if (ptr->lchild == node)
        {
            ptr->balancefactor++;
        }
        else
        {
            ptr->balancefactor--;
        }

        if (ptr->balancefactor == 0)
        {
            /* 插入后还是平衡树 */
            break;
        }
        else if (ptr->balancefactor == 1 || ptr->balancefactor == -1)
        {
            /* 当前子树平衡，继续向上遍历 */
            node = ptr;
            continue;
        }

        /* 不平衡，需要旋转 */
        if (ptr->balancefactor < 0)           //新节点在ptr的右侧
        {
            if (node->balancefactor < 0)
            {
                /* RR型
                 *       A               C     
                 *      / \             / \
                 *     B   C    ->     A   Ex   
                 *        / \         / \
                 *       D   Ex      B   D
                 */
                rotate_RR(tree, ptr);
            }
            else   
            {
                /* RL型
                 *       A             A                   Dx      
                 *      / \           / \                /    \
                 *     B   C    ->   B   Dx     ->      A      C
                 *        / \           / \            / \    / \
                 *       Dx  E         Dl  C          B   Dl Dr  E
                 *                        / \
                 *                       Dr  E
                 */
                rotate_RL(tree, ptr);
            }
        }
        else                                //新节点在ptr的左侧
        {
            if (node->balancefactor < 0)
            {
                /* LR型
                 *       A               A                 Ex                  
                 *      / \             / \              /    \
                 *     B   C    ->     Ex  C    ->      B      A
                 *    / \             / \              / \    / \
                 *   D   Ex          B   Er           D   El Er  C
                 *                  / \
                 *                 D   El                                          
                 */
                rotate_LR(tree, ptr);
            }
            else
            {
                /* LL型
                 *       A             B
                 *      / \           / \
                 *     B   C    ->   Dx  A 
                 *    / \               / \
                 *   Dx  E             E   C
                 */
                rotate_LL(tree, ptr);
            }
        }
        break;
    }

    return 0;
}

/* 插入 */

int BiTreeInsert(struct bitree_t *tree, struct bitree_node_t *newnode)
{
    if (NULL == tree || NULL == newnode || NULL == tree->cmp)
    {
        return -1;
    }

    newnode->lchild = NULL;
    newnode->rchild = NULL;
    newnode->parent = NULL;
    newnode->balancefactor = 0;

    /* 空树，直接作根 */
    if (NULL == tree->root)
    {
        tree->root = newnode;
        tree->count++;
        return 0;
    }

    /* 树深度不超过2层时，简单插入 */
    if (tree->cmp(newnode, tree->root) < 0)
    {
        if (tree->root->lchild == NULL)
        {
            tree->root->lchild = newnode;
            newnode->parent = tree->root;
            tree->root->balancefactor++;
            tree->count++;
        }
    }
    else
    {
        if (tree->root->rchild == NULL)
        {
            tree->root->rchild = newnode;
            newnode->parent = tree->root;
            tree->root->balancefactor--;
            tree->count++;
        }
    }

    /* AVL平衡二叉树 */
    insert_avl(tree, newnode);
    tree->count++;
    return 0;
}

int BiTreeAddData(struct bitree_t *tree, void *data)
{
    struct bitree_node_t *node = NULL;

    if (NULL == tree || NULL == data)
    {
        return -1;
    }

    node = (struct bitree_node_t *)malloc(sizeof(struct bitree_node_t));
    if (NULL == node)
    {
        return -1;
    }

    node->data = data;
    return BiTreeInsert(tree, node);
}