iOS 红黑树（平衡二叉树）的学习以及演示Demo

2019年4月17日 306次阅读来源: 晨寂

之前有段时间再研究算法，然后学到了平衡二叉树 – 红黑树的理论，于是在纸上画了很多很多，导致一本笔记被我画完了……
想想，要不就干脆开发一个Demo来演示二叉树，省纸！

这是项目的github地址：https://github.com/WalkingToTheDistant/RBTree，如果觉得我写的代码对你有帮助的话，给我一个星呗^ _ ^
可能项目还存在问题（自己测试能力有限），如果你们发现了问题，告诉我一声，我一定尽快修改过来的哈。

首先看看Demo的实际效果

《iOS 红黑树（平衡二叉树）的学习以及演示Demo》添加树节点的操作示例

效果截图

《iOS 红黑树（平衡二叉树）的学习以及演示Demo》红黑树效果截图

首先介绍一下红黑树的规则

所有节点是红色或黑色。
根节点是黑色。
所有叶子节点都是黑色。
从任一节点到其每个叶子节点的所有路径都包含相同数目的黑色节点。
从每个叶子到根的所有路径上不能有两个连续的红色节点。

接下来介绍一下红黑树的平衡旋转方法
首先针对当前节点是父节点的左孩纸的情况讨论，当前节点是右孩纸的话，旋转就相反

情况1：当前节点的兄弟节点为红色，则根据从任意节点出发，各个分子的黑高度是一样的规则，那么兄弟节点的孩纸树中必有黑节点，则把兄弟节点变为黑，父节点变为红，然后以父节点左旋，旋转之后，兄弟节点的左孩纸则变成当前节点的兄弟节点（变成情况2）。
情况2：兄弟节点为黑，其子孩纸都为黑，那么则把兄弟节点变为红，然后当前节点移到其父节点。
情况3：兄弟节点为黑，其左孩纸为红，右孩纸为黑，那么则把左孩纸变黑，兄弟节点变红，然后在兄弟节点右旋，这样当前节点的兄弟节点则变成了右转后的左节点（变成情况4）
情况4：兄弟节点为黑，其右孩纸为红，那么则把兄弟节点的颜色变成其父节点的颜色，而父节点的颜色变成黑，兄弟节点的右孩纸变为红，在父节点左旋
情况5：兄弟节点为空，则直接把当前节点指向其父节点再处理

我开发Demo的时候，大部分时间耗在这个平衡旋转的算法上，痛苦…
额，算法的规则细节，就不说了呗，具体红黑树的逻辑你们可以百度看看，只是如果要理解的话，需要深入模拟哈。

接下来看一下重要的功能代码：
先看下节点的代码结构

typedef enum : int{
    NodeColor_Red = 0,  /* 红色 */
    NodeColor_Black,    /* 黑色 */
    
} NodeColor;

/** 红黑树节点 */
@interface RBTreeNode : NSObject

/** 初始化对象，并复制节点值和节点颜色 */
+ (instancetype) treeNodeWithValue:(NSNumber*)value withColor:(NodeColor)color;

/** 节点颜色 */
@property(nonatomic, assign, setter=setNodeColor:, getter=getNodeColor) NodeColor mNodeColor;

/** 节点值 */
@property(nonatomic, retain, setter=setNodeValue:, getter=getNodeValue) NSNumber *mNodeValue;

/** 左节点 */
@property(nonatomic, retain, setter=setLeftNode:, getter=getLeftNode) RBTreeNode *mLeftNode;

/** 右节点 */
@property(nonatomic, retain, setter=setRightNode:, getter=getRightNode) RBTreeNode *mRightNode;

/** 父节点，设置为Weak避免循环引用的问题 */
@property(nonatomic, weak, setter=setParentNode:, getter=getParentNode) RBTreeNode *mParentNode;

@end

平衡旋转的代码

- (void) handleRBTreeForDeleteNode:(RBTreeNode*)parentNodeForCurrentNode
                   withCurrentNode:(RBTreeNode*)currentNode
                   withBrotherNode:(RBTreeNode*)brotherForCurrentNode
                 withIsLeftOrRight:(BOOL)isLeftOfParentNodeForCurrentNode
{
    /**
     下面针对当前节点是父节点的左孩纸的情况讨论，当前节点是右孩纸的话，旋转就相反
     情况1：当前节点的兄弟节点为红色，则根据从任意节点出发，各个分子的黑高度是一样的规则，那么兄弟节点的孩纸树中必有黑节点，则把兄弟节点变为黑，父节点变为红，然后以父节点左旋，旋转之后，兄弟节点的左孩纸则变成当前节点的兄弟节点（变成情况2）。
     情况2：兄弟节点为黑，其子孩纸都为黑，那么则把兄弟节点变为红，然后当前节点移到其父节点。
     情况3：兄弟节点为黑，其左孩纸为红，右孩纸为黑，那么则把左孩纸变黑，兄弟节点变红，然后在兄弟节点右旋，这样当前节点的兄弟节点则变成了右转后的左节点（变成情况4）
     情况4：兄弟节点为黑，其右孩纸为红，那么则把兄弟节点的颜色变成其父节点的颜色，而父节点的颜色变成黑，兄弟节点的右孩纸变为红，在父节点左旋
     情况5：兄弟节点为空，则直接把当前节点指向其父节点再处理
     */
    if(currentNode != mRootTreeNode
       && currentNode.getNodeColor == NodeColor_Black)
    {
        RBTreeNode *grandFatherNode = parentNodeForCurrentNode.getParentNode; // 不是根节点，则必存在父节点
        RBTreeNode *uncleNode = nil;
        BOOL isLeftOrRightForGrandFather = YES;
        if(grandFatherNode == nil){
            isLeftOrRightForGrandFather = YES;
            uncleNode = nil;
            
        } else if(grandFatherNode.getLeftNode == parentNodeForCurrentNode){
            isLeftOrRightForGrandFather = YES;
            uncleNode = grandFatherNode.getRightNode;
            
        } else {
            isLeftOrRightForGrandFather = NO;
            uncleNode = grandFatherNode.getLeftNode;
        }
        
        // ======= 先处理没有旋转的情况 =======
        if(brotherForCurrentNode == nil){ // 情况5
            [self handleRBTreeForDeleteNode:grandFatherNode withCurrentNode:parentNodeForCurrentNode withBrotherNode:uncleNode withIsLeftOrRight:isLeftOrRightForGrandFather];
            
        } else if(brotherForCurrentNode.getNodeColor == NodeColor_Black
                  && ( brotherForCurrentNode.getLeftNode == nil || brotherForCurrentNode.getLeftNode.getNodeColor == NodeColor_Black )
                  && ( brotherForCurrentNode.getRightNode == nil || brotherForCurrentNode.getRightNode.getNodeColor == NodeColor_Black) )
        {   // 情况2
            [brotherForCurrentNode setNodeColor:NodeColor_Red];
            [self handleRBTreeForDeleteNode:grandFatherNode withCurrentNode:parentNodeForCurrentNode withBrotherNode:uncleNode withIsLeftOrRight:isLeftOrRightForGrandFather];
            
        }
        // ======= 接下来处理需要进行旋转的情况 =======
        else if(isLeftOfParentNodeForCurrentNode == YES){ // 当前节点是左孩纸：接下来的情况需要进行旋转， 则需要判断当前节点是左孩纸还是右孩纸
            if(brotherForCurrentNode.getNodeColor == NodeColor_Red){ // 情况1
                [brotherForCurrentNode setNodeColor:NodeColor_Black];
                [parentNodeForCurrentNode setNodeColor:NodeColor_Red];
                [self nodeRotate_toLeft:parentNodeForCurrentNode];
                [self handleRBTreeForDeleteNode:parentNodeForCurrentNode withCurrentNode:currentNode withBrotherNode:parentNodeForCurrentNode.getRightNode withIsLeftOrRight:isLeftOfParentNodeForCurrentNode];
         
            } else if(brotherForCurrentNode.getNodeColor == NodeColor_Black
                      && (brotherForCurrentNode.getLeftNode != nil && brotherForCurrentNode.getLeftNode.getNodeColor == NodeColor_Red)
                      && (brotherForCurrentNode.getRightNode == nil || brotherForCurrentNode.getRightNode.getNodeColor == NodeColor_Black))
            {   // 情况3
                [brotherForCurrentNode.getLeftNode setNodeColor:NodeColor_Black];
                [brotherForCurrentNode setNodeColor:NodeColor_Red];
                [self nodeRotate_toRight:brotherForCurrentNode];
                [self handleRBTreeForDeleteNode:parentNodeForCurrentNode
                                withCurrentNode:currentNode
                                withBrotherNode:parentNodeForCurrentNode.getRightNode
                              withIsLeftOrRight:isLeftOfParentNodeForCurrentNode];
                
            } else if(brotherForCurrentNode.getNodeColor == NodeColor_Black
                      && (brotherForCurrentNode.getRightNode != nil && brotherForCurrentNode.getRightNode.getNodeColor == NodeColor_Red))
            {   // 情况4
                [brotherForCurrentNode setNodeColor:[parentNodeForCurrentNode getNodeColor]];
                [parentNodeForCurrentNode setNodeColor:NodeColor_Black];
                [brotherForCurrentNode.getRightNode setNodeColor:NodeColor_Black];
                [self nodeRotate_toLeft:parentNodeForCurrentNode];
                [self handleRBTreeForDeleteNode:parentNodeForCurrentNode
                                withCurrentNode:currentNode
                                withBrotherNode:parentNodeForCurrentNode.getRightNode
                              withIsLeftOrRight:YES];
            }
        
        } else if(isLeftOfParentNodeForCurrentNode != YES){ // 当前节点是右孩纸
            if(brotherForCurrentNode.getNodeColor == NodeColor_Red){ // 情况1
                [brotherForCurrentNode setNodeColor:NodeColor_Black];
                [parentNodeForCurrentNode setNodeColor:NodeColor_Red];
                [self nodeRotate_toRight:parentNodeForCurrentNode];
                [self handleRBTreeForDeleteNode:parentNodeForCurrentNode withCurrentNode:currentNode withBrotherNode:parentNodeForCurrentNode.getLeftNode withIsLeftOrRight:isLeftOfParentNodeForCurrentNode];
                
            } else if(brotherForCurrentNode.getNodeColor == NodeColor_Black
                      && (brotherForCurrentNode.getRightNode != nil && brotherForCurrentNode.getRightNode.getNodeColor == NodeColor_Red)
                      && (brotherForCurrentNode.getLeftNode == nil || brotherForCurrentNode.getLeftNode.getNodeColor == NodeColor_Black))
            {   // 情况3
                [brotherForCurrentNode.getRightNode setNodeColor:NodeColor_Black];
                [brotherForCurrentNode setNodeColor:NodeColor_Red];
                [self nodeRotate_toLeft:brotherForCurrentNode];
                [self handleRBTreeForDeleteNode:parentNodeForCurrentNode
                                withCurrentNode:currentNode
                                withBrotherNode:parentNodeForCurrentNode.getLeftNode
                              withIsLeftOrRight:isLeftOfParentNodeForCurrentNode];
                
            } else if(brotherForCurrentNode.getNodeColor == NodeColor_Black
                      && (brotherForCurrentNode.getLeftNode != nil && brotherForCurrentNode.getLeftNode.getNodeColor == NodeColor_Red))
            {   // 情况4
                [brotherForCurrentNode setNodeColor:[parentNodeForCurrentNode getNodeColor]];
                [parentNodeForCurrentNode setNodeColor:NodeColor_Black];
                [brotherForCurrentNode.getLeftNode setNodeColor:NodeColor_Black];
                [self nodeRotate_toRight:parentNodeForCurrentNode];
                [self handleRBTreeForDeleteNode:parentNodeForCurrentNode
                                withCurrentNode:currentNode
                                withBrotherNode:parentNodeForCurrentNode.getLeftNode
                              withIsLeftOrRight:isLeftOfParentNodeForCurrentNode];
            }
            
        }
    }
    [currentNode setNodeColor:NodeColor_Black];
}

添加新节点时，判断新节点的位置，以及旋转红黑树的操作

/** 根据当前指向的节点调整红黑树 */
- (void) handleRBTreeForAddNode:(RBTreeNode*)treeNode
{
    /** 情况1：红黑树只有一个根节点 */
    if(treeNode == mRootTreeNode){ // 如果是根节点，那么直接将节点变黑即可
        [treeNode setNodeColor:NodeColor_Black];
        return;
    }
    /** 以下情况则必定存在父节点 */
    /** 情况2：当前节点的父节点为黑色 */
    if(treeNode.getParentNode.getNodeColor == NodeColor_Black){ // 符合红黑树规则，不需要处理
        return;
    }
    
    /** 以下情况父节点必定是红色，所以按照红黑树规则，必定存在祖父节点（规则：红黑树的根节点必是黑色） */
    RBTreeNode *parentNode = treeNode.getParentNode; // 父节点
    RBTreeNode *grandfatherNode = parentNode.getParentNode; // 祖父节点
    BOOL isLeftForParentNodeToGrandfather = (grandfatherNode.getLeftNode == parentNode)? YES : NO; // 父节点是左节点还是右节点，这个决定处理红黑树旋转的方向
    RBTreeNode *uncleNode = (isLeftForParentNodeToGrandfather == YES)? grandfatherNode.getRightNode : grandfatherNode.getLeftNode; // 叔叔节点`
    BOOL isLeftForCurrentNodeToParent = (parentNode.getLeftNode == treeNode)? YES : NO; // 当前节点是左节点还是右节点
    
    /** 情况3：当前节点的父节点为红色，而叔叔节点也为红色（祖父节点和叔叔节点不为空） */
    if(parentNode.getNodeColor == NodeColor_Red
        && uncleNode != nil
        && uncleNode.getNodeColor == NodeColor_Red){
        // 把父节点、叔叔节点都变为黑，祖父节点变为红，然后当前节点指针指向祖父节点再继续监测红黑树
        [parentNode setNodeColor:NodeColor_Black];
        [uncleNode setNodeColor:NodeColor_Black];
        [grandfatherNode setNodeColor:NodeColor_Red];
        return [self handleRBTreeForAddNode:grandfatherNode];
    }
    
    /** 情况4：当前节点的父节点为红色，而叔叔节点为黑色，或者为空 */
    if(parentNode.getNodeColor == NodeColor_Red
        && (uncleNode == nil || uncleNode.getNodeColor == NodeColor_Black)){
        
        if(isLeftForParentNodeToGrandfather == YES){ // 父节点是祖父节点的左孩纸
            if(isLeftForCurrentNodeToParent == YES){ // 当前节点是父节点的左孩纸，这时候只需要把父节点变为黑，祖父节点变成红，当前节点指向祖父节点，并在祖父节点向右旋转
                [parentNode setNodeColor:NodeColor_Black];
                [grandfatherNode setNodeColor:NodeColor_Red];
                [self nodeRotate_toRight:grandfatherNode];
                return [self handleRBTreeForAddNode:grandfatherNode];
            } else { // 当前节点是父节点的右孩纸，那么先把当前节点指向父节点，然后在父节点左旋
                [self nodeRotate_toLeft:parentNode];
                return [self handleRBTreeForAddNode:parentNode];
            }
            
        } else { // 父节点是祖父节点的右孩纸
            if(isLeftForCurrentNodeToParent == YES){// 当前节点是父节点的左孩纸，那么先把当前节点指向父节点，然后在父节点右旋
                [self nodeRotate_toRight:parentNode];
                return [self handleRBTreeForAddNode:parentNode];
            
            } else {// 当前节点是父节点的右孩纸，这时候只需要把父节点变为黑，祖父节点变成红，当前节点指向祖父节点，并在祖父节点向左旋转
                [parentNode setNodeColor:NodeColor_Black];
                [grandfatherNode setNodeColor:NodeColor_Red];
                [self nodeRotate_toLeft:grandfatherNode];
                return [self handleRBTreeForAddNode:grandfatherNode];
            }
        }
    }
    [mRootTreeNode setParentNode:nil];
}

    原文作者：晨寂
    原文地址: https://www.jianshu.com/p/2e7ef5753bb1
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。