【数据结构】AVL树代码(Java版)

AVL树

• AVL树
• BinaryTree.java
• BST.java
• AVLTree.java

AVL树是在二叉搜索树的基础上学习的，【数据结构】二叉搜索树

AVL树

• AVL 树是晶早发明的自平衡二叉搜索树之—
• 平衡因子（Balance Factor）：某结点的左右子树的高度差
• AVL树的特点
  • 每个节点的平衡因子只可能是1、0、-1
    (绝对值≤1，如果超过1，称之为“失衡")
  • 每个节点的左右子树高度差不超过1
  • 搜索、添加、删除的时间复杂度是$o(logn)$

BinaryTree.java

在二叉搜索树的 BinaryTree.java 的基础上增加了 createNode() 方法，原因是 AVL树 拥有自己的节点（增加了height属性）。

/**
 * 二叉树
 */
@SuppressWarnings("unchecked")
public class BinaryTree<E> {
	protected int size; // 元素数量
	protected Node<E> root; // 根节点
	
	/**
	 * 访问器接口 ——> 访问器抽象类
	 * 增强遍历接口
	 */
	public static abstract class Visitor<E> {
		boolean stop;
		// 如果返回true，就代表停止遍历
		abstract boolean visit(E element);
	}
	/**
	 * 内部类，节点类
	 */
	protected static class Node<E> {
		E element; // 元素值
		Node<E> left; // 左节点
		Node<E> right; // 右节点
		Node<E> parent; // 父节点

		public Node(E element, Node<E> parent) {
			this.element = element;
			this.parent = parent;
		}
		public boolean isLeaf(){ // 是否叶子节点
			return left==null && right==null;
		}
		public boolean hasTwoChildren(){ // 是否有两个子节点
			return left!=null && right!=null;
		}
		public boolean isLeftChild(){ // 判断自己是不是左子树
			return parent!=null && this==parent.left;
		}
		public boolean isRightChild(){ // 判断自己是不是右子树
			return parent!=null && this==parent.right;
		}
	}

	/**
	 * 创建节点的方法，用于给AVL树创建节点
	 */
	protected Node<E> createNode(E element, Node<E> parent){
		return new Node<>(element, parent); // 默认返回一个通用节点
	}
	
}

BST.java

/**
 * 二叉搜索树
 */
@SuppressWarnings("unchecked")
public class BST<E> extends BinaryTree<E> {
	
	// 比较器，根据传入的比较器实现 compareTo() 方法
	private Comparator<E> comparator;
	
	public BST (Comparator<E> comparator){ // 可以传一个比较器
		this.comparator = comparator;
	}
	public BST(){ // 不传比较器，相当于传入一个 null
		this(null); //
	}
	
	/**
	 * 添加元素
	 */
	public void add(E element) {
		elementNotNullCheck(element); // 不能传入空节点
		// 传入第一个节点
		if(root == null){
			root = createNode(element, null);
			size++;
			// 新添加节点之后的处理
			afterAdd(root);
			return;
		}
		// 添加的不是第一个节点
		// 找到父节点
		Node<E> parent = root;
		Node<E> node = root;
		int cmp = 0;
		do {
			cmp = compareTo(node.element, element); // 方向
			parent = node; // 父节点
			if(cmp < 0){
				node = node.right;
			}else if(cmp > 0){
				node = node.left;
			}else{ // 相等，最好是覆盖掉，也可以采取其他操作，看具体需求
				node.element = element;
				return;
			}
		} while (node != null);
		// 插入到父节点的哪个位置
		Node<E> newNode = createNode(element, parent);
		if(cmp < 0){
			parent.right = newNode;
		}else{
			parent.left = newNode;
		}
		size++;
		
		// 新添加节点之后的处理
		afterAdd(newNode);
	}
	/**
	 * 根据传入的值删除元素
	 */
	public void remove(E element) {
		remove(node(element));
	}
	// 根据节点删除元素
	private void remove(Node<E> node) {
		if (node == null) return;
		
		size--;
		
		if (node.hasTwoChildren()) { // 度为2的节点
			// 找到后继节点
			Node<E> s = successor(node);
			// 用后继节点的值覆盖度为2的节点的值
			node.element = s.element;
			// 删除后继节点
			node = s;
		}
		
		// 删除node节点（node的度必然是1或者0）
		Node<E> replacement = node.left != null ? node.left : node.right;
		
		if (replacement != null) { // node是度为1的节点
			// 更改parent
			replacement.parent = node.parent;
			// 更改parent的left、right的指向
			if (node.parent == null) { // node是度为1的节点并且是根节点
				root = replacement;
			} else if (node == node.parent.left) {
				node.parent.left = replacement;
			} else { // node == node.parent.right
				node.parent.right = replacement;
			}
			
			// 删除节点后的调整
			afterRemove(node);
		} else if (node.parent == null) { // node是叶子节点并且是根节点
			root = null;
			
			// 删除节点后的调整
			afterRemove(node);
		} else { // node是叶子节点，但不是根节点
			if (node == node.parent.left) {
				node.parent.left = null;
			} else { // node == node.parent.right
				node.parent.right = null;
			}
			
			// 删除节点后的调整
			afterRemove(node);
		}
	}
	/**
	 * 添加node之后的调整，用于给AVL树继承后重写
	 * 在BST中无需调整，但是在AVL树中需要
	 */
	protected void afterAdd(Node<E> node){}
	/**
	 * 删除node之后的调整，用于给AVL树继承后重写
	 * 在BST中无需调整，但是在AVL树中需要
	 */
	protected void afterRemove(Node<E> node){}
	
	// 根据元素值获取节点元素
	private Node<E> node(E element){
		elementNotNullCheck(element);
		
		Node<E> node = root;
		while(node != null){
			int cmp = compareTo(element, node.element);
			if(cmp < 0){
				node = node.left;
			}else if (cmp > 0){
				node = node.right;
			}else{ // cmp == 0
				return node;
			}
		}
		return null;
	}
	// 节点元素比较
	private int compareTo(E e1, E e2) {
		if (comparator != null) { // 传入比较器则通过比较器比较
			return comparator.compare(e1, e2);
		}
		// 没传比较器，元素内部必须自行实现了 Comparable 接口
		return ((Comparable<E>)e1).compareTo(e2);
	}
	// 检测传入的节点是否为空
	private void elementNotNullCheck(E element) {
		if (element == null) { // 不能传入空节点
			throw new IllegalArgumentException("element must not null");
		}
	}
	
}

AVLTree.java

AVL 树继承二叉搜索树。

public class AVLTree<E> extends BST<E> {
	
	// AVL树的节点，需要计算平衡因子，因此比普通二叉树多维护一个height属性
	// (将height放入普通二叉树里没有用处，浪费空间)
	private static class AVLNode<E> extends Node<E>{
		int height = 1;
		
		public AVLNode(E element, Node<E> parent) {
			super(element, parent);
		}
		public int balanceFactor(){ // 获取该节点平衡因子
			int leftHeight = left==null ? 0 : ((AVLNode<E>)left).height;
			int rightHeight = right==null ? 0 : ((AVLNode<E>)right).height;
			return leftHeight - rightHeight;
		}
		public void updateHeight(){ // 更新高度
			int leftHeight = left==null ? 0 : ((AVLNode<E>)left).height;
			int rightHeight = right==null ? 0 : ((AVLNode<E>)right).height;
			height = 1 + Math.max(leftHeight, rightHeight);
		}
		public Node<E> tallerChild(){
			int leftHeight = left==null ? 0 : ((AVLNode<E>)left).height;
			int rightHeight = right==null ? 0 : ((AVLNode<E>)right).height;
			if(leftHeight > rightHeight) return left;
			if(rightHeight > leftHeight) return right;
			// 高度一样则返回同方向的，左子节点则返回左，否则返回右
			return isLeftChild() ? left : right;
		}
		@Override
		public String toString() {
			String parentString = "null";
			if(parent != null){
				parentString = parent.element.toString();
			}
			return element + "_p(" + parentString + ")_h(" + height + ")";
		}
	}
	
	public AVLTree(Comparator<E> comparator){
		super(comparator);
	}
	public AVLTree() {
		this(null);
	}
	
	/**
	 * 增加节点后的调整
	 */
	@Override
	protected void afterAdd(Node<E> node) {
		while ((node = node.parent) != null) {
			if(isBalanced(node)){ // 如果平衡
				// 更新高度
				updateHeight(node);
			}else{ // 如果不平衡
				// 恢复平衡
				rebalance(node);
				// 整棵树恢复平衡
				break;
			}
		}
	}
	/**
	 * 删除节点后的调整
	 */
	@Override
	protected void afterRemove(Node<E> node) {
		while ((node = node.parent) != null) {
			if (isBalanced(node)) {
				// 更新高度
				updateHeight(node);
			} else {
				// 恢复平衡
				rebalance(node);
			}
		}
	}
	/**
	 * 重写父类中的 createNode
	 * 返回 AVLNode
	 */
	@Override
	protected Node<E> createNode(E element, Node<E> parent) {
		return new AVLNode<>(element, parent);
	}
	/**
	 * 判断传入节点是否平衡（平衡因子的绝对值 <= 1）
	 */
	private boolean isBalanced(Node<E> node){ 
		return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
	}
	/**
	 * 更新高度
	 */
	private void updateHeight(Node<E> node){
		((AVLNode<E>)node).updateHeight();
	}
	/**
	 * 恢复平衡
	 * @param node 高度最低的那个不平衡节点
	 */
	private void rebalance2(Node<E> grand){
		Node<E> parent = ((AVLNode<E>)grand).tallerChild();
		Node<E> node = ((AVLNode<E>)parent).tallerChild();
		if(parent.isLeftChild()){//L
			if(node.isLeftChild()){//LL
				rotateRight(grand);//LL则右旋
			}else{//LR
				rotateLeft(parent);
				rotateRight(grand);
			}
		}else{//R
			if(node.isLeftChild()){//RL
				rotateRight(parent);
				rotateLeft(grand);
			}else{//RR
				rotateLeft(grand);//RR则左旋
			}
		}
	}
	/**
	 * 恢复平衡
	 * 一种神奇的做法
	 */
	private void rebalance(Node<E> grand){
		Node<E> parent = ((AVLNode<E>)grand).tallerChild();
		Node<E> node = ((AVLNode<E>)parent).tallerChild();
		if(parent.isLeftChild()){//L
			if(node.isLeftChild()){//LL
				rotate(grand, node, node.right, parent, parent.right, grand);
			}else{//LR
				rotate(grand, parent, node.left, node, node.right, grand);
			}
		}else{//R
			if(node.isLeftChild()){//RL
				rotate(grand, grand, node.left, node, node.right, parent);
			}else{//RR
				rotate(grand, grand, parent.left, parent, node.left, node);
			}
		}
	}
	/**
	 * 统一旋转
	 */
	private void rotate(
			Node<E> r, // 子树的根节点
			Node<E> b, Node<E> c,
			Node<E> d,
			Node<E> e, Node<E> f) {
		// 让d成为这颗子树的根结点
		d.parent = r.parent;
		if(r.isLeftChild()){
			r.parent.left = d;
		}else if(r.isRightChild()){
			r.parent.right = d;
		}else{
			root = d;
		}
		// b-c
		b.right = c;
		if(c!=null){
			c.parent = b;
		}
		updateHeight(b);
		
		// e-f
		f.left = e;
		if(e != null){
			e.parent = f;
		}
		updateHeight(f);
		
		// b-d-f
		d.left = b;
		d.right = f;
		b.parent = d;
		f.parent = d;
		updateHeight(d);
	}
	/*private void rotate(
			Node<E> r, // 子树的根节点
			Node<E> a, Node<E> b, Node<E> c,
			Node<E> d,
			Node<E> e, Node<E> f, Node<E> g) {
		// 让d成为这颗子树的根结点
		d.parent = r.parent;
		if(r.isLeftChild()){
			r.parent.left = d;
		}else if(r.isRightChild()){
			r.parent.right = d;
		}else{
			root = d;
		}
		// a-b-c
		b.left = a;
		if(a!=null){
			a.parent = b;
		}
		b.right = c;
		if(c!=null){
			c.parent = b;
		}
		updateHeight(b);
		
		// e-f-g
		f.left = e;
		if(e != null){
			e.parent = f;
		}
		f.right = g;
		if(g != null){
			g.parent = f;
		}
		updateHeight(f);
		
		// b-d-f
		d.left = b;
		d.right = f;
		b.parent = d;
		f.parent = d;
		updateHeight(d);
	}*/
	/**
	 * 左旋转
	 */
	private void rotateLeft(Node<E> grand){
		Node<E> parent = grand.right;
		Node<E> child = parent.left;
		grand.right = child;
		parent.left = grand;
		
		afterRotate(grand, parent, child);
	}
	/**
	 * 右旋转
	 */
	private void rotateRight(Node<E> grand){
		Node<E> parent = grand.left;
		Node<E> child = parent.right;
		grand.left = child;
		parent.right = grand;
		
		afterRotate(grand, parent, child);
	}
	private void afterRotate(Node<E> grand, Node<E> parent, Node<E> child){
		// 让parent成为子树的根节点
			parent.parent = grand.parent;
			if(grand.isLeftChild()){
				grand.parent.left = parent;
			}else if (grand.isRightChild()){
				grand.parent.right = parent;
			}else{// grand是root节点
				root = parent;
			}
			
			// 更新child的parent
			if(child != null){
			child.parent = grand;
			}
			
			// 更新grand的parent
			grand.parent = parent;
			
			// 更新高度
			updateHeight(grand);
			updateHeight(parent);
	}
	
}

来源：CSDN

作者：Jerry Low

链接：https://blog.csdn.net/weixin_43734095/article/details/104728663