AVL樹的java實現
阿新 • • 發佈:2018-12-24
AVL樹類,其中用Node類封裝了元素,左右兒子和高度,以此來作為結點:
public class AVLTree<E extends Comparable<E>> {
// 樹根結點
private Node<E> root;
// 向樹中插入資料
public Node<E> add(E element) {
return root = insert(element, root);
}
// 刪除樹中值為element的元素
public Node<E> delete(E element) {
return root = remove(element, root);
}
// 前序遍歷樹
public void print() {
print(root);
}
/**
* 向樹中插入資料
*
* @param element 資料的值
* @param node 樹的根結點
* @return 返回插入資料後的樹的根結點
*/
private Node<E> insert(E element, Node<E> node) {
if (node == null) {
return new Node<>(element);
}
if (element.compareTo(node.element) < 0) {
node.left = insert(element, node.left);
} else if (element.compareTo(node.element) > 0) {
node.right = insert (element, node.right);
}
calcHeight(node);
return balance(node);
}
/**
* 刪除樹中的元素
*
* @param element 要刪除的元素值
* @param node 樹的根結點
* @return 返回刪除後的樹根結點
*/
private Node<E> remove(E element, Node<E> node) {
if (node == null || (node.left == null && node.right == null)) {
return null;
}
if (element.compareTo(node.element) < 0) {
node.left = remove(element, node.left);
} else if (element.compareTo(node.element) > 0) {
node.right = remove(element, node.right);
} else {
if (node.right == null) {// 右空左不空
node = node.left;
} else if (node.left == null) {// 左空右不空
node = node.right;
} else {//左右都不空,則取出右子樹最小結點,並用來替換根結點
Node<E> rightMin = searchMin(node.right);
node.element = rightMin.element;
node.right = remove(rightMin.element, node.right);
}
}
calcHeight(node);
return balance(node);
}
/**
* 列印以node為樹根的樹
*
* @param node 樹根
*/
private void print(Node<E> node) {
if (node == null) {
return;
}
System.out.println(node.element + " , height = " + node.height);
print(node.left);
print(node.right);
}
/**
* AVL樹的結點類
*
* @param <E>結點值的型別
*/
static class Node<E> {
E element;
Node<E> left;
Node<E> right;
int height;
public Node(E element) {
this.element = element;
}
}
private Node<E> searchMin(Node<E> node) {
assert node != null;
if (node.left != null) {
return searchMin(node.left);
}
return node;
}
/**
* 計算結點的高度
*
* @param node 要計算的節點
* @return 返回節點高度
*/
private int height(Node<E> node) {
return node == null ? -1 : node.height;
}
private void calcHeight(Node<E> node) {
node.height = Math.max(height(node.left), height(node.right)) + 1;
}
/**
* 左旋
*
* @param node 要旋轉的子樹的根結點
* @return 返回旋轉後的子樹的根結點
*/
private Node<E> leftRotate(Node<E> node) {
Node<E> newNode = node.right;
node.right = newNode.left;
newNode.left = node;
calcHeight(node);
calcHeight(newNode);
return newNode;
}
/**
* 右旋
*
* @param node 要旋轉的子樹的根結點
* @return 返回旋轉後的子樹的根結點
*/
private Node<E> rightRotate(Node<E> node) {
Node<E> newNode = node.left;
node.left = newNode.right;
newNode.right = node;
calcHeight(node);
calcHeight(newNode);
return newNode;
}
/**
* 先左旋再右旋
*
* @param node 要旋轉的子樹的根結點
* @return 返回旋轉後的子樹的根結點
*/
private Node<E> leftAndRightRotate(Node<E> node) {
node.left = leftRotate(node.left);
return rightRotate(node);
}
/**
* 先右旋再左旋
*
* @param node 要旋轉的子樹的根結點
* @return 返回旋轉後的子樹的根結點
*/
private Node<E> rightAndLeftRotate(Node<E> node) {
node.right = rightRotate(node.right);
return leftRotate(node);
}
/**
* 讓以node為根結點的樹恢復平衡
*
* @param node 根結點
* @return 返回恢復平衡後的樹的根結點
*/
private Node<E> balance(Node<E> node) {
// assert node != null;
if (height(node.left) - height(node.right) == 2) {
if (height(node.left.left) > height(node.left.right)) {// 需要進行右旋轉
return rightRotate(node);
} else {// 需要左旋再右旋
return leftAndRightRotate(node);
}
} else if (height(node.right) - height(node.left) == 2) {
if (height(node.right.right) > height(node.right.left)) {// 需要進行左旋轉
return leftRotate(node);
} else {// 需要右旋再左旋
return rightAndLeftRotate(node);
}
}
return node;
}
}
測試程式碼:
public class Main {
public static void main(String[] args) {
AVLTree<Integer> tree = new AVLTree<>();
tree.add(3);
tree.add(2);
tree.add(1);
tree.add(4);
tree.add(5);
tree.add(6);
tree.add(7);
tree.add(10);
tree.add(9);
tree.add(8);
tree.print();
System.out.println("=====================");
tree.delete(4);
tree.print();
}
}