programming-examples/java/Data_Structures/TreapImplicitKey.java

228 lines
6.1 KiB
Java
Raw Normal View History

2019-11-15 12:59:38 +01:00
import java.util.*;
// https://en.wikipedia.org/wiki/Treap
public class TreapImplicitKey {
// Modify the following 5 methods to implement your custom operations on the tree.
// This example implements Add/Max operations. Operations like Add/Sum, Set/Max can also be implemented.
static int modifyOperation(int x, int y) {
return x + y;
}
// query (or combine) operation
static int queryOperation(int leftValue, int rightValue) {
return Math.max(leftValue, rightValue);
}
static int deltaEffectOnSegment(int delta, int segmentLength) {
if (delta == getNeutralDelta()) return getNeutralDelta();
// Here you must write a fast equivalent of following slow code:
// int result = delta;
// for (int i = 1; i < segmentLength; i++) result = queryOperation(result, delta);
// return result;
return delta;
}
static int getNeutralDelta() {
return 0;
}
static int getNeutralValue() {
return Integer.MIN_VALUE;
}
// generic code
static Random random = new Random();
static int joinValueWithDelta(int value, int delta) {
if (delta == getNeutralDelta()) return value;
return modifyOperation(value, delta);
}
static int joinDeltas(int delta1, int delta2) {
if (delta1 == getNeutralDelta()) return delta2;
if (delta2 == getNeutralDelta()) return delta1;
return modifyOperation(delta1, delta2);
}
static void pushDelta(Treap root) {
if (root == null)
return;
root.nodeValue = joinValueWithDelta(root.nodeValue, root.delta);
root.subTreeValue = joinValueWithDelta(root.subTreeValue, deltaEffectOnSegment(root.delta, root.size));
if (root.left != null)
root.left.delta = joinDeltas(root.left.delta, root.delta);
if (root.right != null)
root.right.delta = joinDeltas(root.right.delta, root.delta);
root.delta = getNeutralDelta();
}
public static class Treap {
int nodeValue;
int subTreeValue;
int delta; // delta affects nodeValue, subTreeValue, left.delta and right.delta
int size;
long prio;
Treap left;
Treap right;
Treap(int value) {
nodeValue = value;
subTreeValue = value;
delta = getNeutralDelta();
size = 1;
prio = random.nextLong();
}
void update() {
subTreeValue = queryOperation(queryOperation(getSubTreeValue(left), joinValueWithDelta(nodeValue, delta)), getSubTreeValue(right));
size = 1 + getSize(left) + getSize(right);
}
}
static int getSize(Treap root) {
return root == null ? 0 : root.size;
}
static int getSubTreeValue(Treap root) {
return root == null ? getNeutralValue() : joinValueWithDelta(root.subTreeValue, deltaEffectOnSegment(root.delta, root.size));
}
public static class TreapPair {
Treap left;
Treap right;
TreapPair(Treap left, Treap right) {
this.left = left;
this.right = right;
}
}
public static TreapPair split(Treap root, int minRight) {
if (root == null)
return new TreapPair(null, null);
pushDelta(root);
if (getSize(root.left) >= minRight) {
TreapPair sub = split(root.left, minRight);
root.left = sub.right;
root.update();
sub.right = root;
return sub;
} else {
TreapPair sub = split(root.right, minRight - getSize(root.left) - 1);
root.right = sub.left;
root.update();
sub.left = root;
return sub;
}
}
public static Treap merge(Treap left, Treap right) {
pushDelta(left);
pushDelta(right);
if (left == null)
return right;
if (right == null)
return left;
// if (random.nextInt(left.size + right.size) < left.size) {
if (left.prio > right.prio) {
left.right = merge(left.right, right);
left.update();
return left;
} else {
right.left = merge(left, right.left);
right.update();
return right;
}
}
public static Treap insert(Treap root, int index, int value) {
TreapPair t = split(root, index);
return merge(merge(t.left, new Treap(value)), t.right);
}
public static Treap remove(Treap root, int index) {
TreapPair t = split(root, index);
return merge(t.left, split(t.right, index + 1 - getSize(t.left)).right);
}
public static Treap modify(Treap root, int a, int b, int delta) {
TreapPair t1 = split(root, b + 1);
TreapPair t2 = split(t1.left, a);
t2.right.delta = joinDeltas(t2.right.delta, delta);
return merge(merge(t2.left, t2.right), t1.right);
}
public static class TreapAndResult {
Treap treap;
int value;
TreapAndResult(Treap t, int value) {
this.treap = t;
this.value = value;
}
}
public static TreapAndResult query(Treap root, int a, int b) {
TreapPair t1 = split(root, b + 1);
TreapPair t2 = split(t1.left, a);
int value = getSubTreeValue(t2.right);
return new TreapAndResult(merge(merge(t2.left, t2.right), t1.right), value);
}
public static void print(Treap root) {
if (root == null)
return;
pushDelta(root);
print(root.left);
System.out.print(root.nodeValue + " ");
print(root.right);
}
// Random test
public static void main(String[] args) {
Treap treap = null;
List<Integer> list = new ArrayList<>();
Random rnd = new Random(1);
for (int step = 0; step < 100000; step++) {
int cmd = rnd.nextInt(6);
if (cmd < 2 && list.size() < 100) {
int pos = rnd.nextInt(list.size() + 1);
int value = rnd.nextInt(100);
list.add(pos, value);
treap = insert(treap, pos, value);
} else if (cmd < 3 && list.size() > 0) {
int pos = rnd.nextInt(list.size());
list.remove(pos);
treap = remove(treap, pos);
} else if (cmd < 4 && list.size() > 0) {
int b = rnd.nextInt(list.size());
int a = rnd.nextInt(b + 1);
int res = list.get(a);
for (int i = a + 1; i <= b; i++)
res = queryOperation(res, list.get(i));
TreapAndResult tr = query(treap, a, b);
treap = tr.treap;
if (res != tr.value)
throw new RuntimeException();
} else if (cmd < 5 && list.size() > 0) {
int b = rnd.nextInt(list.size());
int a = rnd.nextInt(b + 1);
int delta = rnd.nextInt(100) - 50;
for (int i = a; i <= b; i++)
list.set(i, joinValueWithDelta(list.get(i), delta));
treap = modify(treap, a, b, delta);
} else {
for (int i = 0; i < list.size(); i++) {
TreapAndResult tr = query(treap, i, i);
treap = tr.treap;
int v = tr.value;
if (list.get(i) != v)
throw new RuntimeException();
}
}
}
System.out.println("Test passed");
}
}