@SuppressWarnings("unchecked")
private void forEach(AVLTree t, PSet.Consumer<K> action) {
if (t.isEmpty()) {
return;
}
action.accept((K) t.key());
forEach(t.left(), action);
forEach(t.right(), action);
}
@Test
public void no_change() {
AVLTree<String, String> t0 = AVLTree.create();
AVLTree<String, String> t1 = t0.put("1", "1");
assertThat(t1.put("1", "1")).isSameAs(t1);
assertThat(t1.remove("3")).isSameAs(t1);
AVLTree<String, String> t2 = t0.put("2", "2");
assertThat(t2.put("2", "2")).isSameAs(t2);
assertThat(t2.remove("3")).isSameAs(t2);
}
private static AVLTree combineTrees(AVLTree l, AVLTree r) {
if (l.isEmpty()) {
return r;
}
if (r.isEmpty()) {
return l;
}
NodeRef oldNode = new NodeRef();
AVLTree newRight = removeMinBinding(r, oldNode);
return balance(l, oldNode.node.key(), oldNode.node.value(), newRight);
}
@Override
public int hashCode() {
if (hashCode == 0) {
int result = key.hashCode();
result = result * 31 + left.hashCode();
result = result * 31 + right.hashCode();
result = result * 31 + value.hashCode();
hashCode = result;
}
return hashCode;
}
@Test
public void test() {
List<Integer> keys = new ArrayList<>();
for (int i = 0; i < 100; i++) {
keys.add(i);
}
Collections.shuffle(keys);
AVLTree<Integer, Object> t = AVLTree.create();
for (Integer key : keys) {
t = t.add(key);
assertThat(t.add(key)).isSameAs(t);
}
assertThat(Counter.countSet(t)).isEqualTo(100);
assertThat(Counter.countMap(t)).isEqualTo(100);
assertThat(t.height()).isGreaterThanOrEqualTo(8).isLessThanOrEqualTo(10);
for (Integer key : keys) {
assertThat(t.contains(key)).isTrue();
t = t.remove(key);
assertThat(t.remove(key)).isSameAs(t);
}
assertThat(Counter.countSet(t)).isEqualTo(0);
assertThat(Counter.countMap(t)).isEqualTo(0);
}
@Test
public void test_replace_root() {
AVLTree<String, String> t0 = AVLTree.create();
AVLTree<String, String> t1 = t0.put("1", "a");
AVLTree<String, String> t2 = t1.put("1", "b");
assertThat(t1).isNotSameAs(t2);
assertThat(t1.get("1")).isEqualTo("a");
assertThat(t2.get("1")).isEqualTo("b");
}
private static AVLTree removeMinBinding(AVLTree t, NodeRef noderemoved) {
assert !t.isEmpty();
if (t.left().isEmpty()) {
noderemoved.node = t;
return t.right();
}
return balance(removeMinBinding(t.left(), noderemoved), t.key(), t.value(), t.right());
}
@Test
public void test_empty() {
AVLTree<String, String> t = AVLTree.create();
assertThat(t).as("singleton").isSameAs(AVLTree.create());
assertThat(t.get("anything")).isNull();
assertThat(t.remove("anything")).isSameAs(t);
assertThat(t.toString()).isEqualTo("");
assertThat(t.hashCode()).isEqualTo(0);
}
/**
* Subtraction must not be used for comparison of keys due to possibility of integer overflow,
* this for example will be the case for sequence below, which was generated using random number
* generator.
*/
@Test
public void do_not_use_subtraction_for_comparison_of_keys() {
Key[] keys = {
new Key(2043979982, ""),
new Key(-36348207, ""),
new Key(-1864559204, ""),
new Key(-2018458363, ""),
new Key(-152409201, ""),
new Key(-1786252453, ""),
new Key(-1853960690, "")
};
AVLTree<Object, Object> t = AVLTree.create();
for (Key key : keys) {
t = t.add(key);
}
for (Key key : keys) {
assertThat(t.get(key)).as("found").isNotNull();
assertThat(t.remove(key)).as("removed").isNotSameAs(t);
}
}
@Test
public void balancing_should_preserve_buckets() {
Object k1 = new Key(1, "k1");
Object k2 = new Key(2, "k2");
Object k3 = new Key(3, "k3");
Object k4 = new Key(4, "k4");
AVLTree<Object, Object> t = AVLTree.create().put(k1, "v1").put(k2, "v2").put(k3, "v3");
Object k1_1 = new Key(1, "k1_1");
t = t.put(k1_1, "v1_1");
t = t.put(k4, "v4");
assertThat(t.height()).as("height after balancing").isEqualTo(3);
assertThat(t.get(k1_1)).isEqualTo("v1_1");
}
@SuppressWarnings("unchecked")
@Nullable
@Override
public V get(K key) {
Preconditions.checkNotNull(key);
AVLTree t = this;
while (!t.isEmpty()) {
int c = KEY_COMPARATOR.compare(key, t.key());
if (c == 0) {
return (V) t.value();
} else if (c < 0) {
t = t.left();
} else {
t = t.right();
}
}
return null;
}
private static AVLTree remove(Object key, AVLTree t) {
if (t.isEmpty()) {
return t;
}
int result = KEY_COMPARATOR.compare(key, t.key());
if (result == 0) {
return combineTrees(t.left(), t.right());
} else if (result < 0) {
AVLTree left = remove(key, t.left());
if (left == t.left()) {
return t;
}
return balance(left, t.key(), t.value(), t.right());
} else {
AVLTree right = remove(key, t.right());
if (right == t.right()) {
return t;
}
return balance(t.left(), t.key(), t.value(), right);
}
}
@Override
public String toString() {
return " " + key + "->" + value + left.toString() + right.toString();
}
private static int incrementHeight(AVLTree l, AVLTree r) {
return (l.height() > r.height() ? l.height() : r.height()) + 1;
}
@Test
public void hashCode_and_equals_should_not_depend_on_order_of_construction() {
Object o1 = new Key(21, "o1");
Object o2 = new Key(45, "o2");
AVLTree<Object, Object> t1 = AVLTree.create().add(o1).add(o2);
AVLTree<Object, Object> t2 = AVLTree.create().add(o2).add(o1);
assertThat(t1.key()).as("shape is different").isNotEqualTo(t2.key());
assertThat(t1.hashCode()).isEqualTo(t2.hashCode());
assertThat(t1).isEqualTo(t2);
assertThat(t2).isEqualTo(t1);
Object o3 = new Key(0, "o3");
AVLTree<Object, Object> t3 = t1.add(o3);
assertThat(t1.hashCode()).isEqualTo(t3.hashCode());
assertThat(t1).isNotEqualTo(t3);
assertThat(t3).isNotEqualTo(t1);
}
private static AVLTree createNode(AVLTree newLeft, AVLTree oldTree, AVLTree newRight) {
return createNode(newLeft, oldTree.key(), oldTree.value(), newRight);
}
private static AVLTree balance(AVLTree l, Object key, Object value, AVLTree r) {
if (l.height() > r.height() + 2) {
assert !l.isEmpty();
AVLTree ll = l.left();
AVLTree lr = l.right();
if (ll.height() >= lr.height()) {
return createNode(ll, l, createNode(lr, key, value, r));
}
assert !lr.isEmpty();
AVLTree lrl = lr.left();
AVLTree lrr = lr.right();
return createNode(createNode(ll, l, lrl), lr, createNode(lrr, key, value, r));
}
if (r.height() > l.height() + 2) {
assert !r.isEmpty();
AVLTree rl = r.left();
AVLTree rr = r.right();
if (rr.height() >= rl.height()) {
return createNode(createNode(l, key, value, rl), r, rr);
}
assert !rl.isEmpty();
AVLTree rll = rl.left();
AVLTree rlr = rl.right();
return createNode(createNode(l, key, value, rll), rl, createNode(rlr, r, rr));
}
return createNode(l, key, value, r);
}
@Test
public void test_one_element() {
AVLTree<String, String> t0 = AVLTree.create();
AVLTree<String, String> t1 = t0.put("1", "a");
AVLTree<String, String> t2 = t0.put("2", "b");
assertThat(t0).isNotSameAs(t1).isNotSameAs(t2);
assertThat(t1).isNotSameAs(t2);
assertThat(t0.get("1")).isNull();
assertThat(t0.get("2")).isNull();
assertThat(t0.get("3")).isNull();
assertThat(t1.get("1")).isEqualTo("a");
assertThat(t1.get("2")).isNull();
assertThat(t2.get("3")).isNull();
assertThat(t2.get("1")).isNull();
assertThat(t2.get("2")).isEqualTo("b");
assertThat(t2.get("3")).isNull();
}
@Test
public void buckets() {
Object k1 = new Key(42, "k1");
Object k2 = new Key(42, "k2");
Object k3 = new Key(42, "k3");
AVLTree<Object, Object> t = AVLTree.create().put(k1, "v1").put(k2, "v2");
assertThat(t.toString()).as("should create bucket").isEqualTo(" k2->v2 k1->v1");
AVLTree<Object, Object> t2 = AVLTree.create().put(k2, "v2").put(k1, "v1");
assertThat(t2.toString())
.as("toString depends on order of operations")
.isEqualTo(" k1->v1 k2->v2");
assertThat(t.equals(t2)).as("should compare buckets").isTrue();
assertThat(t2.equals(t)).as("should compare buckets").isTrue();
assertThat(t.hashCode())
.isEqualTo(
((31 * k1.hashCode()) ^ "v1".hashCode()) + ((31 * k2.hashCode()) ^ "v2".hashCode()));
assertThat(t2.hashCode())
.as("hashCode doesn't depend on order of operations")
.isEqualTo(t.hashCode());
assertThat(t.get(k1)).isEqualTo("v1");
assertThat(t.get(k2)).isEqualTo("v2");
assertThat(t.get(k3)).as("not such key").isNull();
assertThat(t.put(k2, "new v2").toString())
.as("should replace head of bucket")
.isEqualTo(" k2->new v2 k1->v1");
assertThat(t.put(k1, "new v1").toString())
.as("should replace element of bucket")
.isEqualTo(" k1->new v1 k2->v2");
assertThat(t.put(k1, "v1")).as("should not change").isSameAs(t);
assertThat(t.put(k2, "v2")).as("should not change").isSameAs(t);
assertThat(t.put(k3, "v3").toString())
.as("should add to bucket")
.isEqualTo(" k3->v3 k2->v2 k1->v1");
assertThat(t.remove(k2).toString()).as("should remove head of bucket").isEqualTo(" k1->v1");
assertThat(t.remove(k1).toString()).as("should remove element of bucket").isEqualTo(" k2->v2");
assertThat(t.remove(k1).remove(k2).toString()).as("should remove bucket").isEqualTo("");
assertThat(t.remove(k3)).as("should not change").isSameAs(t);
HashMap<Object, Object> biConsumer = new HashMap<>();
t.forEach((k, v) -> assertThat(biConsumer.put(k, v)).as("unique key-value").isNull());
assertThat(biConsumer).isEqualTo(ImmutableMap.of(k1, "v1", k2, "v2"));
HashSet<Object> consumer = new HashSet<>();
t.forEach(k -> assertThat(consumer.add(k)).as("unique key").isTrue());
assertThat(consumer).containsOnly(k1, k2);
}
private static AVLTree put(Object key, Object value, AVLTree t) {
if (t.isEmpty()) {
return createNode(t, key, value, t);
}
int result = KEY_COMPARATOR.compare(key, t.key());
if (result == 0) {
if (value.equals(t.value())) {
return t;
}
return createNode(t.left(), key, value, t.right());
} else if (result < 0) {
AVLTree left = put(key, value, t.left());
if (left == t.left()) {
return t;
}
return balance(left, t.key(), t.value(), t.right());
} else {
AVLTree right = put(key, value, t.right());
if (right == t.right()) {
return t;
}
return balance(t.left(), t.key(), t.value(), right);
}
}
