Java数据结构--平衡二叉树

DEMO地址：
https://github.com/zhaopingfu/MDataStruct/blob/master/src/com/pf/%E6%A0%91/AVLBintrayTree.java

在这里有一些资源，辅助看的:
https://github.com/zhaopingfu/MDataStruct/tree/master/resources/%E6%A0%91/AVL%E6%A0%91

• 平衡二叉树是一种二叉排序树，其中每一个节点的左子树和右子树的高度差至多等于1.

• 平衡因子：二叉树上节点的左子树深度减去右子树的深度的值成为平衡因子BF（Balance Factor）。

• 最小不平衡子树：距离插入节点最近的，且平衡因子的绝对值大于1的节点为根的子树为最小不平衡树。

• 构建AVL树

  

  /**
   * @author zhaopf
   * @version 1.0
   * @QQ: 1308108803
   * @date 2017年12月25日
   * 平衡二叉树
   */
  public class AVLBintrayTree<E extends Comparable<E>> {
      /**
       * 根节点
       */
      private Node<E> root;
      /**
       * 树的节点
       */
      private int size;
      /**
       * 左边高，左边减去右边
       */
      private static final int LEFT_HEIGHT = 1;
      /**
       * 两边一样高，左边减去右边
       */
      private static final int EQUALS_HEIGHT = 0;
      /**
       * 右边高，左边减去右边
       */
      private static final int RIGHT_HEIGHT = -1;
  
      /**
       * 添加一个元素
       *
       * @param element
       * @return
       */
      public boolean put(E element) {
          if (element == null) {
              return false;
          }
          Node<E> t = root;
          if (t == null) {
              root = new Node<>(element, null);
              size++;
              return true;
          }
          // 开始添加元素
          Node<E> parent = null;
          Comparable<? super E> e = element;
          int compare;
          // 查找父节点
          do {
              parent = t;
              compare = e.compareTo(t.element);
              if (compare > 0) {
                  t = t.right;
              } else if (compare < 0) {
                  t = t.left;
              } else {
                  return false;
              }
          } while (t != null);
          // 添加到父节点上
          Node<E> newNode = new Node<>(element, parent);
          compare = newNode.element.compareTo(parent.element);
          if (compare > 0) {
              parent.right = newNode;
          } else if (compare < 0) {
              parent.left = newNode;
          }
          // 检查平衡问题
          while (parent != null) {
              compare = newNode.element.compareTo(parent.element);
              // 如果节点是添加到右子树
              if (compare > 0) {
                  parent.balance--;
              } else {
                  parent.balance++;
              }
              // 还是平衡的
              if (parent.balance == 0) {
                  break;
              }
              // 如果出现了平衡问题
              if (Math.abs(parent.balance) == 2) {
                  fixAfterInsertion(parent);
                  break;
              } else {
                  // 回溯到父节点
                  parent = parent.parent;
              }
          }
          size++;
          return true;
      }
  
      private void fixAfterInsertion(Node<E> parent) {
          if (parent.balance == 2) {
              leftBalance(parent);
          } else if (parent.balance == -2) {
              rightBalance(parent);
          }
      }
  
      /**
       * 中序遍历节点
       *
       * @return
       */
      public List<Node<E>> midSelectTree() {
          List<Node<E>> result = new ArrayList<>(size);
          Node<E> temp = root;
          if (temp == null) {
              return result;
          }
          Stack<Node<E>> stack = new Stack<>();
          // 添加根节点
          stack.add(temp);
          Set<Node<E>> set = new HashSet<>(size);
          while (!stack.isEmpty()) {
              Node<E> t = stack.pop();
              if (t.getLeft() != null && !set.contains(t.getLeft())) {
                  if (t.getRight() != null) {
                      stack.push(t.getRight());
                  }
                  stack.push(t);
                  stack.push(t.getLeft());
              } else {
                  set.add(t);
                  result.add(t);
                  if (t.getRight() != null && !stack.contains(t.getRight())) {
                      stack.push(t.getRight());
                  }
              }
          }
          return result;
      }
  
      /**
       * 左平衡操作，节点t的不平衡是因为左子树过深
       *
       * @param t
       */
      private void leftBalance(Node<E> t) {
          if (t == null || t.left == null) {
              return;
          }
          Node<E> tLeftChild = t.left;
          switch (tLeftChild.balance) {
              // 如果新的结点插入到t的左孩子的左子树中，则直接进行右旋操作即可
              case LEFT_HEIGHT:
                  t.balance = EQUALS_HEIGHT;
                  tLeftChild.balance = EQUALS_HEIGHT;
                  // 右旋
                  rightRotate(t);
                  break;
              case EQUALS_HEIGHT:
                  break;
              // 如果新的结点插入到t的左孩子的右子树中，则需要进行分情况讨论
              case RIGHT_HEIGHT:
                  Node<E> tLeftChildRightChild = tLeftChild.right;
                  switch (tLeftChildRightChild.balance) {
                      // 当t的左孩子的右子树根节点的balance = LEFT_HIGH
                      case LEFT_HEIGHT:
                          t.balance = RIGHT_HEIGHT;
                          tLeftChild.balance = EQUALS_HEIGHT;
                          tLeftChildRightChild.balance = EQUALS_HEIGHT;
                          break;
                      // 当t的左孩子的右子树根节点的balance = EQUAL_HIGH
                      case EQUALS_HEIGHT:
                          t.balance = EQUALS_HEIGHT;
                          tLeftChild.balance = EQUALS_HEIGHT;
                          tLeftChildRightChild.balance = EQUALS_HEIGHT;
                          break;
                      // 当t的左孩子的右子树根节点的balance = RIGHT_HIGH
                      case RIGHT_HEIGHT:
                          t.balance = EQUALS_HEIGHT;
                          tLeftChild.balance = LEFT_HEIGHT;
                          tLeftChildRightChild.balance = EQUALS_HEIGHT;
                          break;
                      default:
                          break;
                  }
                  // 先将t的左子树左旋
                  leftRotate(tLeftChild);
                  // 在将t右旋
                  rightRotate(t);
                  break;
              default:
                  break;
          }
      }
  
      /**
       * 右平衡操作，即结点t的不平衡是因为右子树过深
       *
       * @param t
       */
      private void rightBalance(Node<E> t) {
          if (t == null || t.right == null) {
              return;
          }
          Node<E> tRightChild = t.right;
          switch (tRightChild.balance) {
              // 如果新的结点插入到p的右孩子的左子树中，则需要进行分情况讨论
              case LEFT_HEIGHT:
                  Node<E> tRightChildLeftChild = tRightChild.left;
                  switch (tRightChildLeftChild.balance) {
                      case LEFT_HEIGHT:
                          t.balance = EQUALS_HEIGHT;
                          tRightChild.balance = RIGHT_HEIGHT;
                          tRightChildLeftChild.balance = EQUALS_HEIGHT;
                          break;
                      case EQUALS_HEIGHT:
                          t.balance = EQUALS_HEIGHT;
                          tRightChild.balance = EQUALS_HEIGHT;
                          tRightChildLeftChild.balance = EQUALS_HEIGHT;
                          break;
                      case RIGHT_HEIGHT:
                          t.balance = LEFT_HEIGHT;
                          tRightChild.balance = EQUALS_HEIGHT;
                          tRightChildLeftChild.balance = EQUALS_HEIGHT;
                          break;
                      default:
                          break;
                  }
                  // 先对t的右孩子进行右旋
                  rightRotate(tRightChild);
                  // 再对t进行左旋
                  leftRotate(t);
                  break;
              case EQUALS_HEIGHT:
                  break;
              // 如果新的结点插入到t的右孩子的右子树中，则直接进行左旋操作即可
              case RIGHT_HEIGHT:
                  t.balance = EQUALS_HEIGHT;
                  tRightChild.balance = EQUALS_HEIGHT;
                  leftRotate(t);
                  break;
              default:
                  break;
          }
      }
  
      /**
       * 左旋
       *
       * @param x 要旋转的节点
       */
      private void leftRotate(Node<E> x) {
          if (x == null || x.right == null) {
              return;
          }
          Node<E> y = x.right;
  
          // step 1
          x.right = y.left;
          if (y.left != null) {
              y.left.parent = x;
          }
  
          // step 2
          if (x.parent == null) {
              root = y;
          } else if (x.parent.left == x) {
              x.parent.left = y;
          } else if (x.parent.right == x) {
              x.parent.right = y;
          }
          y.parent = x.parent;
  
          // step 3
          y.left = x;
          x.parent = y;
      }
  
      /**
       * 右旋
       *
       * @param y 要旋转的节点
       */
      private void rightRotate(Node<E> y) {
          if (y == null || y.left == null) {
              return;
          }
          Node<E> x = y.left;
  
          // step 1
          y.left = x.right;
          if (x.right != null) {
              x.right.parent = y;
          }
  
          // step 2
          if (y.parent == null) {
              root = x;
          } else if (y.parent.left == y) {
              y.parent.left = x;
          } else if (y.parent.right == y) {
              y.parent.right = x;
          }
          x.parent = y.parent;
  
          // step 3
          x.right = y;
          y.parent = x;
      }
  
      public int getSize() {
          return size;
      }
  
      public Node<E> getRoot() {
          return root;
      }
  
      /**
       * 节点
       *
       * @param <E> 实现了Comparable接口
       */
      public class Node<E extends Comparable<E>> {
          /**
           * 数据
           */
          private E element;
          /**
           * 平衡因子
           */
          private int balance;
          /**
           * 左子树
           */
          private Node<E> left;
          /**
           * 右子树
           */
          private Node<E> right;
          /**
           * 父节点
           */
          private Node<E> parent;
  
          public Node(E element, Node<E> parent) {
              this.element = element;
              this.parent = parent;
          }
  
          public E getElement() {
              return element;
          }
  
          public void setElement(E element) {
              this.element = element;
          }
  
          public int getBalance() {
              return balance;
          }
  
          public void setBalance(int balance) {
              this.balance = balance;
          }
  
          public Node<E> getLeft() {
              return left;
          }
  
          public void setLeft(Node<E> left) {
              this.left = left;
          }
  
          public Node<E> getRight() {
              return right;
          }
  
          public void setRight(Node<E> right) {
              this.right = right;
          }
  
          public Node<E> getParent() {
              return parent;
          }
  
          public void setParent(Node<E> parent) {
              this.parent = parent;
          }
      }
  }