平衡二叉树的实现

来自http://zhouyunan2010.iteye.com/blog/1255303

package com.utils;

import java.util.Iterator;
import java.util.NoSuchElementException;

public class AVLTree<E> {
 
 private Entry<E> root = null;
 
 
 private int size = 0;
 
 public AVLTree(){
  
 }
 
 public int size(){
  return size;
 }

 
 private void rotateLeft(Entry<E> p) {
  System.out.println("绕"+p.element+"左旋");
  if(p!=null){
   Entry<E> r = p.right;
   p.right = r.left; //把p右子树的左节点嫁接到p的右节点，如上图，把BL作为A的右子节点
   if (r.left != null) //如果B的左节点BL不为空，把BL的父节点设为A
    r.left.parent = p;
   r.parent = p.parent;  //A的父节点设为B的父节点
   if (p.parent == null)    //如果p是根节点
                root = r;      //r变为父节点，即B为父节点
            else if (p.parent.left == p)  //如果p是左子节点
                p.parent.left = r;        //p的父节点的左子树为r
            else         //如果p是右子节点
                p.parent.right = r;    //p的父节点的右子树为r
            r.left = p;       //p变为r的左子树，即A为B的左子树
            p.parent = r;      //同时更改p的父节点为r，即A的父节点为B
  }
 }
 
 
 private void rotateRight(Entry<E> p){
  System.out.println("绕"+p.element+"右旋");
  if(p!=null){
   Entry<E> l = p.left; 
   p.left = l.right;    //把B的右节点BR作为A的左节点
            if (l.right != null)   //如果BR不为null，设置BR的父节点为A
             l.right.parent = p;
            l.parent = p.parent;  //A的父节点赋给B的父节点
            if (p.parent == null) //如果p是根节点
                root = l;          //B为根节点
            else if (p.parent.right == p) //如果A是其父节点的左子节点
                p.parent.right = l;  //B为A的父节点的左子树
            else       //如果A是其父节点的右子节点
             p.parent.left = l;  //B为A的父节点的右子树
            l.right = p;    //A为B的右子树
            p.parent = l;      //设置A的父节点为B
  }
 }

 
 
 public boolean add(E element){
  Entry<E> t = root;
  if(t == null){
   root = new Entry<E>(element,null);
   size = 1;
   return true;
  }
  int cmp;
  Entry<E> parent; //保存t的父节点
  Comparable<? super E> e = (Comparable<? super E>) element;
  //从根节点向下搜索，找到插入位置
  do {
   parent = t;  
   cmp = e.compareTo(t.element);
   if(cmp < 0){
    t = t.left;
   }else if(cmp > 0){
    t = t.right;
   }else{
    return false;
   }
  } while (t!=null);
  
  Entry<E> child = new Entry(element,parent);
  if(cmp < 0){
   parent.left = child;
   
  }else{
   parent.right = child; 
  }
  //自下向上回溯，查找最近不平衡节点
  while(parent!=null){
   cmp = e.compareTo(parent.element);
   if(cmp < 0){    //插入节点在parent的左子树中
    parent.balance++;
   }else{    //插入节点在parent的右子树中
    parent.balance--;
   }
   if(parent.balance == 0){ //此节点的balance为0，不再向上调整BF值，且不需要旋转
    break;
   }
   if(Math.abs(parent.balance) == 2){ //找到最小不平衡子树根节点
    fixAfterInsertion(parent);
    break;     //不用继续向上回溯
   }
   parent = parent.parent;
  }
  size ++;
  return true;
 }
 
 
 private void fixAfterInsertion(Entry<E> p){
  if(p.balance == 2){
   leftBalance(p);
  }
  if(p.balance == -2){
   rightBalance(p);
  }
 }
 
 
 
 private boolean leftBalance(Entry<E> t){
  boolean heightLower = true;
  Entry<E> l = t.left;
  switch (l.balance) {
  case LH:   //左高，右旋调整,旋转后树的高度减小
   t.balance = l.balance = EH;
   rotateRight(t);
   break;
  case RH:         //右高，分情况调整           
   Entry<E> rd = l.right;
   switch (rd.balance) { //调整各个节点的BF
   case LH: //情况1
    t.balance = RH;
    l.balance = EH;
    break;
   case EH: //情况2
    t.balance = l.balance = EH;
    break;
   case RH: //情况3
    t.balance = EH;
    l.balance = LH;
    break;
   }
   rd.balance = EH;
   rotateLeft(t.left);
   rotateRight(t);
   break;
  case EH:   //特殊情况4,这种情况在添加时不可能出现，只在移除时可能出现，旋转之后整体树高不变
   l.balance = RH;
   t.balance = LH;
   rotateRight(t);
   heightLower = false;
   break;
  }
  return heightLower;
 }
 
 
 private boolean rightBalance(Entry<E> t){
  boolean heightLower = true;
  Entry<E> r = t.right;
  switch (r.balance) {
  case LH:   //左高，分情况调整
   Entry<E> ld = r.left;
   switch (ld.balance) { //调整各个节点的BF
   case LH: //情况1
    t.balance = EH;
    r.balance = RH;
    break;
   case EH:    //情况2
    t.balance = r.balance = EH;
    break;
   case RH: //情况3
    t.balance = LH;
    r.balance = EH;
    break;
   }
   ld.balance = EH;
   rotateRight(t.right);
   rotateLeft(t);
   break;
  case RH:   //右高，左旋调整
   t.balance = r.balance = EH;
   rotateLeft(t);
   break;
  case EH:       //特殊情况4
   r.balance = LH;
   t.balance = RH;
   rotateLeft(t);
   heightLower = false;
   break;
  }
  return heightLower;
 }
 
 
 private Entry<E> getEntry(Object element) { 
     Entry<E> tmp = root; 
     Comparable<? super E> e = (Comparable<? super E>) element;
     int c; 
     while (tmp != null) { 
         c = e.compareTo(tmp.element); 
         if (c == 0) { 
             return tmp; 
         } else if (c < 0) { 
             tmp = tmp.left; 
         } else { 
             tmp = tmp.right; 
         } 
     } 
     return null; 
 } 
 
 
 public boolean remove(Object o){
  Entry<E> e = getEntry(o);
  if(e!=null){
   deleteEntry(e);
   return true;
  }
  return false;
 }
 
 private void deleteEntry(Entry<E> p){
  size --;
  //如果p左右子树都不为空，找到其直接后继，替换p，之后p指向s，删除p其实是删除s
  //所有的删除左右子树不为空的节点都可以调整为删除左右子树有其一不为空，或都为空的情况。
  if (p.left != null && p.right != null) {
    Entry<E> s = successor(p);
    p.element = s.element;
    p = s;
  }
  Entry<E> replacement = (p.left != null ? p.left : p.right);

        if (replacement != null) {  //如果其左右子树有其一不为空
            replacement.parent = p.parent;
            if (p.parent == null) //如果p为root节点
                root = replacement;
            else if (p == p.parent.left) //如果p是左孩子
                p.parent.left  = replacement; 
            else       //如果p是右孩子
                p.parent.right = replacement;

            p.left = p.right = p.parent = null;  //p的指针清空
           
            //这里更改了replacement的父节点，所以可以直接从它开始向上回溯
            fixAfterDeletion(replacement); 

        } else if (p.parent == null) { // 如果全树只有一个节点
            root = null;
        } else {  //左右子树都为空
         fixAfterDeletion(p); //这里从p开始回溯
            if (p.parent != null) {
                if (p == p.parent.left)
                    p.parent.left = null;
                else if (p == p.parent.right)
                    p.parent.right = null;
                p.parent = null;
            }
        } 
 }
 
 
 static <E> Entry<E> successor(Entry<E> t) {
  if (t == null)
   return null;
  else if (t.right != null) { //往右，然后向左直到尽头
   Entry<E> p = t.right;
   while (p.left != null)
    p = p.left;
   return p;
  } else {  //right为空，如果t是p的左子树，则p为t的直接后继
   Entry<E> p = t.parent;
   Entry<E> ch = t;
   while (p != null && ch == p.right) { //如果t是p的右子树，则继续向上搜索其直接后继
    ch = p;
    p = p.parent;
   }
   return p;
  }
    }
 
 
 private void fixAfterDeletion(Entry<E> p){
  boolean heightLower = true;  //看最小子树调整后，它的高度是否发生变化，如果减小，继续回溯
  Entry<E> t = p.parent;
  Comparable<? super E> e = (Comparable<? super E>)p.element;
  int cmp;
  //自下向上回溯，查找不平衡的节点进行调整
  while(t!=null && heightLower){
   cmp = e.compareTo(t.element);
   
   if(cmp >= 0 ){  
    t.balance ++;
   }else{
    t.balance --;
   }
   if(Math.abs(t.balance) == 1){ //父节点经过调整平衡因子后，如果为1或-1，说明调整之前是0，停止回溯。
    break;
   }
   Entry<E> r = t;
   //这里的调整跟插入一样
   if(t.balance == 2){
    heightLower = leftBalance(r);
   }else if(t.balance==-2){
    heightLower = rightBalance(r);
   }
   t = t.parent;
  }
 }
 
 
 
 private static final int LH = 1; //左高
 private static final int EH = 0;    //等高
 private static final int RH = -1; //右高
 
 
 static class Entry<E>{
  E element;
  Entry<E> parent;
  Entry<E> left;
  Entry<E> right;
  int balance = EH; //平衡因子，只能为1或0或-1
  
  public Entry(E element,Entry<E> parent){
   this.element = element;
   this.parent = parent;
  }
  
  public Entry(){}

  @Override
  public String toString() {
   return element+" BF="+balance;
  }
   
 }
 
 
 //返回中序遍历此树的迭代器,返回的是一个有序列表
 private class BinarySortIterator implements Iterator<E>{
  Entry<E> next;
     Entry<E> lastReturned;
    
     public BinarySortIterator(){
      Entry<E> s = root;
      if(s !=null){
       while(s.left != null){ //找到中序遍历的第一个元素
        s = s.left;
       }
      }
      next = s; //赋给next
     }
    
  @Override
  public boolean hasNext() {
   return next!=null;
  }

  @Override
  public E next() {
   Entry<E> e = next;
   if (e == null)
    throw new NoSuchElementException();
   next = successor(e);
   lastReturned = e;
   return e.element;
  }

  @Override
  public void remove() {
            if (lastReturned == null)
                throw new IllegalStateException();
            // deleted entries are replaced by their successors
            if (lastReturned.left != null && lastReturned.right != null)
                next = lastReturned;
            deleteEntry(lastReturned);
            lastReturned = null;
  }
 }
 
 public Iterator<E> itrator(){
  return new BinarySortIterator();
 }
 

 
 private int treeHeight(Entry<E> p){
  if(p == null){
   return -1;
  }
  return 1 + Math.max(treeHeight(p.left), treeHeight(p.right));
 }
 
 
 //测试，你也可以任意测试
 public static void main(String[] args) {
  AVLTree<Integer> tree = new AVLTree<Integer>();
  System.out.println("------添加------");
  tree.add(50);
  System.out.print(50+" ");
  tree.add(66);
  System.out.print(66+" ");
  for(int i=0;i<10;i++){
   int ran = (int)(Math.random() * 100);
   System.out.print(ran+" ");
   tree.add(ran);
  }
  System.out.println("------删除------");
  tree.remove(50);
  tree.remove(66);
  
  System.out.println();
  Iterator<Integer> it = tree.itrator();
  while(it.hasNext()){
   System.out.print(it.next()+" ");
  }
 }
}