Skip to content

Instantly share code, notes, and snippets.

@yoavst
Last active May 17, 2018 15:35
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save yoavst/ddac8a068c8f00a3547508cd8ecddbec to your computer and use it in GitHub Desktop.
Save yoavst/ddac8a068c8f00a3547508cd8ecddbec to your computer and use it in GitHub Desktop.
WAVL Tree
package il.ac.tau.cs.ds.hw1;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
/**
* Represent an inner (non-virtual) node in a {@link WAVLTree}.
*/
class WAVLInnerNode extends WAVLNode {
private int key;
@NotNull
private String value;
@NotNull
private WAVLNode left;
@NotNull
private WAVLNode right;
@Nullable
private WAVLNode parent;
private int subtreeSize;
private int rank;
WAVLInnerNode(int key, @NotNull String value) {
this(key, value, new WAVLVirtualNode(), new WAVLVirtualNode());
}
WAVLInnerNode(int key, @NotNull String value, @NotNull WAVLNode left, @NotNull WAVLNode right) {
this.key = key;
this.value = value;
this.left = left;
this.right = right;
this.rank = Math.max(left.getRank(), right.getRank()) + 1;
this.subtreeSize = left.getSubtreeSize() + right.getSubtreeSize() + 1;
left.setParent(this);
right.setParent(this);
}
@Override
public void setLeft(@NotNull WAVLNode left) {
this.left = left;
}
@Override
public void setRight(@NotNull WAVLNode right) {
this.right = right;
}
@Override
public void setParent(@Nullable WAVLNode parent) {
this.parent = parent;
}
public void setRank(int rank) {
this.rank = rank;
}
@Override
public int getKey() {
return key;
}
@NotNull
@Override
public String getValue() {
return value;
}
@NotNull
@Override
public WAVLNode getLeft() {
return left;
}
@NotNull
@Override
public WAVLNode getRight() {
return right;
}
@Override
public boolean isInnerNode() {
return true;
}
@Override
public int getSubtreeSize() {
return subtreeSize;
}
@Nullable
@Override
public WAVLNode getParent() {
return parent;
}
@Override
public int getRank() {
return rank;
}
@Override
public String toString() {
return "Node(" + getValue() + ", rank=" + getRank() + ")";
}
void updateSubtreeSize() {
subtreeSize = left.getSubtreeSize() + right.getSubtreeSize() + 1;
}
void updateChildren() {
getLeft().setParent(this);
getRight().setParent(this);
}
}
package il.ac.tau.cs.ds.hw1;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
/**
* Represents a node of {@link WAVLTree} binary tree.
* Can represent a virtual node if {@link #isInnerNode()} returns false.
*/
public abstract class WAVLNode {
static final int EXTERNAL_NODE_RANK = -1;
public abstract int getKey();
@Nullable
public abstract String getValue();
@Nullable
public abstract WAVLNode getLeft();
@Nullable
public abstract WAVLNode getRight();
@Nullable
public abstract WAVLNode getParent();
public abstract boolean isInnerNode();
public abstract int getSubtreeSize();
public abstract int getRank();
public abstract void setLeft(@NotNull WAVLNode left);
public abstract void setRight(@NotNull WAVLNode right);
public abstract void setParent(@Nullable WAVLNode parent);
/**
* @return true if is non-virtual leaf node
*/
boolean isLeaf() {
if (!isInnerNode()) return false;
// Note: cannot be null since inner node
//noinspection ConstantConditions
return !getLeft().isInnerNode() && !getRight().isInnerNode();
}
/**
* @return The rank diff between the node and its left child. Could not be called on virtual node.
*/
int leftRankDiff() {
WAVLNode left = getLeft();
if (left == null)
throw new IllegalStateException("Left node cannot be null");
return getRank() - left.getRank();
}
/**
* @return The rank diff between the node and its right child. Could not be called on virtual node.
*/
int rightRankDiff() {
WAVLNode right = getRight();
if (right == null)
throw new IllegalStateException("Left node cannot be null");
return getRank() - right.getRank();
}
}
package il.ac.tau.cs.ds.hw1;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.function.BiConsumer;
/**
* An implementation of a WAVL Tree.
* (Haupler, Sen & Tarajan ‘15)
*/
@SuppressWarnings("WeakerAccess")
public class WAVLTree {
/**
* The tree root node. Should not be virtual node unless the tree is empty.
*/
@NotNull
private WAVLNode root = new WAVLVirtualNode();
@NotNull
private WAVLNode min = root;
@NotNull
private WAVLNode max = root;
//region Insert operation
/**
* inserts an item with key k and info i to the WAVL tree.
* the tree must remain valid (keep its invariants).
*
* @return the number of rebalancing operations, or 0 if no rebalancing operations were necessary.
* -1 if an item with key k already exists in the tree.
*/
@SuppressWarnings("UnusedReturnValue")
public int insert(int key, String value) {
final WAVLInnerNode node = new WAVLInnerNode(key, value);
// If the tree is empty, insert the node as the root.
if (!root.isInnerNode()) {
root = node;
min = root;
max = root;
return 0;
}
final WAVLNode insertionPlace = insertionPlace(root, key);
if (insertionPlace.isInnerNode()) {
return -1;
} else {
// Note: getParent cannot be null since tree is not empty
@NotNull @SuppressWarnings("ConstantConditions") final WAVLInnerNode parent = (WAVLInnerNode) insertionPlace.getParent();
setChildBySideAndUpdateChild(parent, node, getSide(parent, insertionPlace));
if (node.getKey() < min.getKey())
min = node;
else if (node.getKey() > max.getKey())
max = node;
return insertionReblance(node);
}
}
/**
* Rebalance the tree starting from the inserted node.
*
* @return the number of rebalancing operations, or 0 if no rebalancing operations were necessary.
*/
private int insertionReblance(WAVLInnerNode node) {
int rebalanceSteps = 0;
while (node.getParent() != null) {
final WAVLInnerNode parent = (WAVLInnerNode) node.getParent();
if (parent.getRank() == node.getRank()) {
final boolean otherSide = !getSide(parent, node);
if (getRankDiffBySide(parent, otherSide) == 1) {
// Case 1:
// * k
// * k * k -1
// Solution: Promote
parent.setRank(parent.getRank() + 1);
rebalanceSteps++;
} else {
if (getRankDiffBySide(node, otherSide) == 2) {
// Case 2:
// * k
// * k * k-2
// * k-1 * k-2
// Solution: single rotation
rotate(parent, otherSide);
rebalanceSteps += 2;
} else {
// Case 3:
// * k
// * k * k-2
// * k-2 * k-1
rotate(node, !otherSide);
rotate(parent, otherSide);
// Note: parent's parent is not null because of the double rotation
final WAVLInnerNode promotedNode = (WAVLInnerNode) parent.getParent();
//noinspection ConstantConditions
promotedNode.setRank(promotedNode.getRank() + 1);
rebalanceSteps += 5;
}
// Note: Reblance complete since rotation (or double rotation) is terminal.
// Just need to update parents for new subtree size, so continuing the loop.
}
}
parent.updateSubtreeSize();
node = parent;
}
root = node;
return rebalanceSteps;
}
/**
* Recursively insertionPlace for the key, starting with currentNode.
*
* @return the insert location for an item with the given key
*/
@NotNull
private static WAVLNode insertionPlace(@NotNull WAVLNode currentNode, int key) {
if (!currentNode.isInnerNode())
return currentNode;
if (currentNode.getKey() == key)
return currentNode;
// Note: getLeft() and getRight() cannot be null since it is an inner node.
//noinspection ConstantConditions
return insertionPlace(currentNode.getKey() > key ? currentNode.getLeft() : currentNode.getRight(), key);
}
//endregion
//region Delete operation
/**
* deletes an item with key k from the binary tree, if it is there;
* the tree must remain valid (keep its invariants).
*
* @return the number of rebalancing operations, or 0 if no rebalancing operations were needed. -1 if an item with key k was not found in the tree.
*/
@SuppressWarnings("UnusedReturnValue")
public int delete(int key) {
if (size() == 0) {
return -1;
}
final WAVLNode searchedDeletionNode = insertionPlace(root, key);
if (!searchedDeletionNode.isInnerNode()) {
return -1;
}
final WAVLInnerNode deletionNode = (WAVLInnerNode) searchedDeletionNode;
final WAVLInnerNode parent = (WAVLInnerNode) deletionNode.getParent();
final boolean sideOfParent = parent == null || getSide(parent, deletionNode);
@NotNull WAVLNode reblanceNode;
if (deletionNode.isLeaf()) {
if (parent == null) {
// Tree of 1 item.
root = new WAVLVirtualNode();
min = root;
max = root;
return 0;
} else {
// a leaf
reblanceNode = new WAVLVirtualNode(parent);
setChildBySide(parent, reblanceNode, sideOfParent);
deleteUpdateMinMax(deletionNode);
}
} else {
deleteUpdateMinMax(deletionNode);
// Note: deletion node is an inner node
//noinspection ConstantConditions
if (deletionNode.getLeft().isInnerNode() ^ deletionNode.getRight().isInnerNode()) {
// Unary node
final WAVLInnerNode replacement = (WAVLInnerNode) (deletionNode.getLeft().isInnerNode() ? deletionNode.getLeft() : deletionNode.getRight());
if (parent == null) {
root = replacement;
root.setParent(null);
} else {
setChildBySideAndUpdateChild(parent, replacement, sideOfParent);
}
reblanceNode = replacement;
} else {
// binary node
// We replace it with its successor or predecessor, based on which subtree of the root it is on.
boolean side = deletionNode.getKey() > root.getKey();
WAVLInnerNode replacement = side ? successor(deletionNode) : predecessor(deletionNode);
// Note: parent cannot be null for non root node
//noinspection ConstantConditions
final WAVLInnerNode replacementParent = (WAVLInnerNode) replacement.getParent();
if (replacementParent != deletionNode) {
// Set `side` child to parent with the `!side` child of replacement
reblanceNode = getChildBySide(replacement, side);
//noinspection ConstantConditions
setChildBySideAndUpdateChild(replacementParent, reblanceNode, !side);
} else {
setChildBySide(deletionNode, getChildBySide(replacement, side), side);
reblanceNode = getChildBySide(deletionNode, side);
}
if (parent == null) {
root = replacement;
root.setParent(null);
} else {
setChildBySideAndUpdateChild(parent, replacement, sideOfParent);
}
replacement.setLeft(deletionNode.getLeft());
replacement.setRight(deletionNode.getRight());
replacement.updateChildren();
replacement.setRank(deletionNode.getRank());
}
}
return deletionReblance(reblanceNode);
}
/**
* Rebalance the tree starting from the deleted node.
*
* @return the number of rebalancing operations, or 0 if no rebalancing operations were necessary.
*/
private int deletionReblance(WAVLNode deletedNode) {
WAVLNode node = deletedNode;
int rebalanceSteps = 0;
while (node.getParent() != null) {
if (node.isLeaf()) ((WAVLInnerNode) node).setRank(0);
final WAVLInnerNode parent = (WAVLInnerNode) node.getParent();
final boolean side = getSide(parent, node);
if (getRankDiffBySide(parent, side) == 3) {
if (getRankDiffBySide(parent, !side) == 2) {
// Case 1:
// * k+3
// * k * k+1
// Solution: Demote
parent.setRank(parent.getRank() - 1);
rebalanceSteps++;
} else {
final WAVLInnerNode otherSide = (WAVLInnerNode) getChildBySide(parent, !side);
if (otherSide.leftRankDiff() == 2 && otherSide.rightRankDiff() == 2) {
// Case 2:
// * k+3
// * k * k+2
// * k * k
// Solution: Double demote
parent.setRank(parent.getRank() - 1);
otherSide.setRank(otherSide.getRank() - 1);
rebalanceSteps += 2;
} else if (getRankDiffBySide(otherSide, !side) == 2) {
// case 4:
// * k+3
// * k * k+2
// * k+1 * k
// Solution: Double rotate
rotate(otherSide, !side);
rotate(parent, side);
parent.setRank(parent.getRank() - 1);
// Note: parent's parent cannot be null since rotated
//noinspection ConstantConditions
((WAVLInnerNode) parent.getParent()).setRank(parent.getParent().getRank() + 2);
rebalanceSteps += 5;
if (parent == root) {
// Note: Parent is no longer root so it has parent
//noinspection ConstantConditions
root = parent.getParent();
}
// Note: Reblance complete!
// Just need to update parents for new subtree size, so continuing the loop.
} else {
// case 3:
// * k+3
// * k * k+2
// * k,k+1 * k+1
// Solution: rotate (+ optional demote)
rotate(parent, side);
otherSide.setRank(otherSide.getRank() + 1);
rebalanceSteps += 3;
if (parent.leftRankDiff() == 2 && parent.rightRankDiff() == 2) {
parent.setRank(parent.getRank() - 1);
rebalanceSteps++;
}
if (parent == root) {
root = otherSide;
}
// Note: Reblance complete!
// Just need to update parents for new subtree size, so continuing the loop.
}
}
}
parent.updateSubtreeSize();
node = node.getParent();
}
if (node.isLeaf()) {
// If we had 2 items before, and left with root only, we need to clear its rank.
((WAVLInnerNode) node).setRank(0);
}
return rebalanceSteps;
}
/**
* Update the min and max based on if they were deleted
*/
private void deleteUpdateMinMax(WAVLInnerNode deletionNode) {
if (deletionNode == min) {
// Note: not null since tree has more than one item
//noinspection ConstantConditions
min = successor(deletionNode);
} else if (deletionNode == max) {
// Note: not null since tree has more than one item
//noinspection ConstantConditions
max = predecessor(deletionNode);
}
}
/**
* @return The successor to the given node, or null if it has no successor.
*/
@Nullable
private static WAVLInnerNode successor(@NotNull WAVLInnerNode node) {
if (node.getRight().isInnerNode()) {
// Note: Cannot be null since inner node
//noinspection ConstantConditions
node = (WAVLInnerNode) node.getRight();
//noinspection ConstantConditions
while (node.getLeft().isInnerNode()) {
node = (WAVLInnerNode) node.getLeft();
}
return node;
}
WAVLInnerNode parent = (WAVLInnerNode) node.getParent();
while (parent != null && node == parent.getRight()) {
node = parent;
parent = (WAVLInnerNode) node.getParent();
}
return parent;
}
/**
* @return The predecessor to the given node, or null if it has no predecessor.
*/
@Nullable
private static WAVLInnerNode predecessor(@NotNull WAVLInnerNode node) {
if (node.getLeft().isInnerNode()) {
// Note: Cannot be null since inner node
//noinspection ConstantConditions
node = (WAVLInnerNode) node.getLeft();
//noinspection ConstantConditions
while (node.getRight().isInnerNode()) {
node = (WAVLInnerNode) node.getRight();
}
return node;
}
WAVLInnerNode parent = (WAVLInnerNode) node.getParent();
while (parent != null && node == parent.getLeft()) {
node = parent;
parent = (WAVLInnerNode) node.getParent();
}
return parent;
}
//endregion
/**
* Rotate parent's [rotateToSide] child with parent where true means rotate the left child to be the new parent.
*
* @param parent the node who is going to be rotated with its children
* @param rotateToSide the side to rotate to
* @see #getSide(WAVLNode, WAVLNode)
*/
private static void rotate(WAVLInnerNode parent, boolean rotateToSide) {
// assuming rotateToSide is right (true)
// * parent * node
// * node * Y ==> * X * parent
// * X * oldOther * oldOther * Y
final WAVLInnerNode node = (WAVLInnerNode) getChildBySide(parent, !rotateToSide);
final WAVLNode oldOther = getChildBySide(node, rotateToSide);
// move oldOther under parent
setChildBySideAndUpdateChild(parent, oldOther, !rotateToSide);
// connect node to parent's parent
node.setParent(parent.getParent());
if (parent.getParent() != null) {
setChildBySide((WAVLInnerNode) parent.getParent(), node, getSide(parent.getParent(), parent));
}
// set parent as child of node
setChildBySideAndUpdateChild(node, parent, rotateToSide);
parent.setRank(parent.getRank() - 1);
// update subtree sizes
parent.updateSubtreeSize();
node.updateSubtreeSize();
}
/**
* Example 1: select(1) returns the value of the node with minimal key
* Example 2: select(size()) returns the value of the node with maximal key
* Example 3: select(2) returns the value 2nd smallest minimal node, i.e the value of the node minimal node's successor
*
* @return the value of the i'th smallest key (return -1 if tree is empty)
*/
@SuppressWarnings("ConstantConditions")
//FIXME @NotNull
@Nullable
public String select(int index) {
if (empty() || (index < 1 || index > size()))
//FIXME replace with return "-1";
return null;
int currentIndex = root.getLeft().getSubtreeSize() + 1;
@NotNull
WAVLNode node = root;
while (currentIndex != index) {
if (currentIndex > index) {
node = node.getLeft();
currentIndex -= node.getRight().getSubtreeSize() + 1;
} else {
node = node.getRight();
currentIndex += node.getLeft().getSubtreeSize() + 1;
}
}
return node.getValue();
}
//region Simple required ADT methods
/**
* @return true if and only if the tree is empty
*/
public boolean empty() {
return size() == 0;
}
/**
* @return the number of nodes in the tree.
*/
public int size() {
return root.getSubtreeSize();
}
/**
* @return the root WAVL node, or null if the tree is empty
*/
@NotNull
public WAVLNode getRoot() {
return root;
}
/**
* @return a sorted array which contains all keys in the tree, or an empty array if the tree is empty.
*/
@NotNull
public int[] keysToArray() {
final int[] arr = new int[size()];
inOrderPass(root, 0, (node, index) -> arr[index] = node.getKey());
return arr;
}
/**
* @return an array which contains all info in the tree, sorted by their respective keys, or an empty array if the tree is empty.
*/
@NotNull
public String[] infoToArray() {
final String[] arr = new String[size()];
inOrderPass(root, 0, (node, index) -> arr[index] = node.getValue());
return arr;
}
/**
* @return the info of the item with the smallest key in the tree,or null if the tree is empty
*/
@Nullable
public String min() {
return empty() ? null : min.getValue();
}
/**
* @return the info of the item with the largest key in the tree, or null if the tree is empty
*/
@Nullable
public String max() {
return empty() ? null : max.getValue();
}
/**
* @return the info of an item with key k if it exists in the tree otherwise null
*/
@Nullable
public String search(int key) {
final WAVLNode node = insertionPlace(root, key);
if (!node.isInnerNode())
return null;
return node.getValue();
}
//endregion
//region Utils
/**
* Goes over the WAVL tree with the given node as a parent, and applying callback on each node with its index, starting from 0
*
* @param node the root of the tree
* @param index the current index, default should be 0
* @param callback what to do on each node based on itself and its in-order scan index
* @return the node's index after consuming all the nodes under it
*/
@SuppressWarnings("ConstantConditions")
private static int inOrderPass(@NotNull WAVLNode node, int index, @NotNull BiConsumer<WAVLNode, Integer> callback) {
if (!node.isInnerNode())
return index;
// Note: the following calls are safe since node is inner node.
final int newIndex = inOrderPass(node.getLeft(), index, callback);
callback.accept(node, newIndex);
return inOrderPass(node.getRight(), newIndex + 1, callback);
}
/**
* Note: node must be a direct child of parent
*
* @return the side of node compared to parent, where true means right and false means left.
*/
private static boolean getSide(@NotNull WAVLNode parent, @NotNull WAVLNode node) {
return parent.getRight() == node;
}
/**
* @return The child at the given child
* @see #getSide(WAVLNode, WAVLNode)
*/
@NotNull
private static WAVLNode getChildBySide(@NotNull WAVLInnerNode parent, boolean side) {
return side ? parent.getRight() : parent.getLeft();
}
/**
* set the parent's child on side.
*
* @see #getSide(WAVLNode, WAVLNode)
*/
private static void setChildBySide(@NotNull WAVLInnerNode parent, @NotNull WAVLNode child, boolean side) {
if (side) {
parent.setRight(child);
} else {
parent.setLeft(child);
}
}
/**
* set the parent's child on side.
*
* @see #getSide(WAVLNode, WAVLNode)
*/
private static void setChildBySideAndUpdateChild(@NotNull WAVLInnerNode parent, @NotNull WAVLNode child, boolean side) {
if (side) {
parent.setRight(child);
} else {
parent.setLeft(child);
}
parent.updateChildren();
}
/**
* @return The rank diff between the parent and the child on it's side.
*/
private static int getRankDiffBySide(@NotNull WAVLNode parent, boolean side) {
return side ? parent.rightRankDiff() : parent.leftRankDiff();
}
@Override
public String toString() {
return WAVLTreePrinter.toString(this);
}
//endregion
}
package il.ac.tau.cs.ds.hw1;
import org.jetbrains.annotations.NotNull;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
class WAVLTreePrinter {
private WAVLTreePrinter() {
}
@NotNull
static String toString(@NotNull WAVLTree tree) {
return toString(tree, true);
}
@NotNull
private static String toString(@NotNull WAVLTree tree, boolean byKey) {
return String.join(System.lineSeparator(), representation(tree.getRoot(), byKey));
}
@SuppressWarnings("ConstantConditions")
@NotNull
private static List<String> representation(@NotNull WAVLNode node, boolean byKey) {
if (!node.isInnerNode()) {
return Collections.singletonList("#");
}
String thisString = byKey ? String.valueOf(node.getKey()) : node.getValue();
return concatenation(representation(node.getLeft(), byKey), thisString, representation(node.getRight(), byKey));
}
@NotNull
private static List<String> concatenation(@NotNull List<String> left, @NotNull String root, @NotNull List<String> right) {
int lwid = left.get(left.size() - 1).length();
int rwid = right.get(right.size() - 1).length();
int rootwid = root.length();
ArrayList<String> result = new ArrayList<>();
result.add(repeat(lwid + 1, " ") + root + repeat(rwid + 1, " "));
int ls = leftSpace(left.get(0));
int rs = rightSpace(right.get(0));
result.add(repeat(ls, " ") + repeat(lwid - ls, "_") + "/" + repeat(rootwid, " ") + "\\" + repeat(rs, "_") + repeat(rwid - rs, " "));
for (int i = 0; i < Math.max(left.size(), right.size()); i++) {
String row = "";
if (i < left.size()) {
row += left.get(i);
} else {
row += repeat(lwid, " ");
}
row += repeat(rootwid + 2, " ");
if (i < right.size()) {
row += right.get(i);
} else {
row += repeat(rwid, " ");
}
result.add(row);
}
return result;
}
@NotNull
private static String repeat(int n, @NotNull String s) {
if (n <= 0) {
return "";
}
StringBuilder builder = new StringBuilder(s.length() * n);
for (int i = 0; i < n; i++) {
builder.append(s);
}
return builder.toString();
}
private static int leftSpace(@NotNull String row) {
int i = row.length() - 1;
while (row.charAt(i) == ' ') {
i -= 1;
}
return i + 1;
}
private static int rightSpace(@NotNull String row) {
int i = 0;
while (row.charAt(i) == ' ') {
i += 1;
}
return i;
}
}
package il.ac.tau.cs.ds.hw1;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
class WAVLVirtualNode extends WAVLNode {
@Nullable
private WAVLNode parent;
public WAVLVirtualNode() {
this(null);
}
public WAVLVirtualNode(@Nullable WAVLNode parent) {
this.parent = parent;
}
@Override
public int getKey() {
return EXTERNAL_NODE_RANK;
}
@Nullable
@Override
public String getValue() {
return null;
}
@Nullable
@Override
public WAVLNode getLeft() {
return null;
}
@Nullable
@Override
public WAVLNode getRight() {
return null;
}
@Nullable
@Override
public WAVLNode getParent() {
return parent;
}
@Override
public boolean isInnerNode() {
return false;
}
@Override
public int getSubtreeSize() {
return 0;
}
@Override
public int getRank() {
return EXTERNAL_NODE_RANK;
}
@Override
public void setLeft(@NotNull WAVLNode left) {
throw new UnsupportedOperationException("Should not call setLeft for a virtual node");
}
@Override
public void setRight(@NotNull WAVLNode right) {
throw new UnsupportedOperationException("Should not call setRight for a virtual node");
}
@Override
public void setParent(@Nullable WAVLNode parent) {
this.parent = parent;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment