diff --git a/pom.xml b/pom.xml
index 9021c41..2d3c032 100644
--- a/pom.xml
+++ b/pom.xml
@@ -43,6 +43,16 @@
3.2.0
test
+
+ org.hamcrest
+ hamcrest
+ 2.1
+
+
+ org.mockito
+ mockito-core
+ 2.27.0
+
diff --git a/src/main/java/com/lodborg/intervaltree/TreeNode.java b/src/main/java/com/lodborg/intervaltree/TreeNode.java
index c40c435..8053acd 100644
--- a/src/main/java/com/lodborg/intervaltree/TreeNode.java
+++ b/src/main/java/com/lodborg/intervaltree/TreeNode.java
@@ -196,6 +196,7 @@ private TreeNode balanceOut(){
// The tree is right-heavy.
if (height(right.left) > height(right.right)){
this.right = this.right.rightRotate();
+ this.right.height = Math.max(height(this.right.left), height(this.right.right)) + 1;
return leftRotate();
} else{
return leftRotate();
@@ -204,6 +205,7 @@ private TreeNode balanceOut(){
// The tree is left-heavy.
if (height(left.right) > height(left.left)){
this.left = this.left.leftRotate();
+ this.left.height = Math.max(height(this.left.left), height(this.left.right)) + 1;
return rightRotate();
} else
return rightRotate();
diff --git a/src/test/java/com/lodborg/intervaltree/TreeNodeTest.java b/src/test/java/com/lodborg/intervaltree/TreeNodeTest.java
index 3b44e47..6a21a8f 100644
--- a/src/test/java/com/lodborg/intervaltree/TreeNodeTest.java
+++ b/src/test/java/com/lodborg/intervaltree/TreeNodeTest.java
@@ -1,12 +1,16 @@
package com.lodborg.intervaltree;
+import com.lodborg.intervaltree.Interval.Bounded;
+import org.hamcrest.CoreMatchers;
import org.junit.Test;
+import org.mockito.Mockito;
+import java.util.*;
+
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.core.Is.is;
import static org.junit.Assert.*;
-import com.lodborg.intervaltree.Interval.*;
-
-import java.util.*;
public class TreeNodeTest {
@Test
@@ -234,4 +238,87 @@ public void test_iteratorRemoveDeletesInnerNodeAndPromotesTheSubtreeRoot(){
assertTrue(list.contains(next));
}
}
+
+ @Test
+ public void leftRightRotation() {
+
+ IntervalTree tree = Mockito.mock(IntervalTree.class);
+ TreeNode node = TreeNode.addInterval(tree, null, new IntegerInterval(100, 100, Bounded.CLOSED));
+ node = TreeNode.addInterval(tree, node, new IntegerInterval(50, 50, Bounded.CLOSED));
+ node = TreeNode.addInterval(tree, node, new IntegerInterval(75, 75, Bounded.CLOSED));
+
+ assertAVLHeightProperty(node);
+ assertHeight(node);
+ }
+
+ @Test
+ public void rightLeftRotation() {
+
+ IntervalTree tree = Mockito.mock(IntervalTree.class);
+ TreeNode node = TreeNode.addInterval(tree, null, new IntegerInterval(50, 50, Bounded.CLOSED));
+ node = TreeNode.addInterval(tree, node, new IntegerInterval(100, 100, Bounded.CLOSED));
+ node = TreeNode.addInterval(tree, node, new IntegerInterval(75, 75, Bounded.CLOSED));
+
+ assertAVLHeightProperty(node);
+ assertHeight(node);
+ }
+
+
+ @Test
+ public void randomIntervals() {
+
+ IntervalTree tree = Mockito.mock(IntervalTree.class);
+ Random random = new Random();
+
+ TreeNode node = TreeNode.addInterval(tree, null, new IntegerInterval(100, 100, Bounded.CLOSED));
+
+ for (int i = 0; i < 100_000; i++) {
+
+ int num = random.nextInt();
+ node = TreeNode.addInterval(tree, node, new IntegerInterval(num - 1, num + 1, Bounded.CLOSED));
+ }
+
+ assertAVLHeightProperty(node);
+ assertHeight(node);
+ }
+
+ private > int assertHeight(TreeNode node) {
+
+ if (node == null) {
+ return 0;
+ }
+
+ int leftHeight = assertHeight(node.left);
+ int rightHeight = assertHeight(node.right);
+
+ int height = Math.max(leftHeight, rightHeight) + 1;
+
+ assertThat(node.height, is(height));
+
+ return height;
+ }
+
+ private > void assertAVLHeightProperty(TreeNode root) {
+
+ if (root == null) {
+ return;
+ }
+
+ assertThat(root.height, greaterThan(0));
+
+ if (root.left == null && root.right == null) {
+ assertThat(root.height, CoreMatchers.is(1));
+ } else {
+ assertThat("Height of two subtrees should not differ by more than 1",
+ Math.abs(height(root.left) - height(root.right)), lessThanOrEqualTo(1));
+ }
+
+ assertAVLHeightProperty(root.left);
+ assertAVLHeightProperty(root.right);
+ }
+
+ private > int height(TreeNode root) {
+
+ return root == null ? 0 : root.height;
+ }
}