private void push(GroupTree z, Centroid start) { while (z.left != null) { if (start == null || start.compareTo(z.leaf) < 0) { // remember we will have to process the right hand branch later stack.push(z.right); // note that there is no guarantee that z.left has any good data z = z.left; } else { // if the left hand branch doesn't contain start, then no push z = z.right; } } // put the leaf value on the stack if it is valid if (start == null || z.leaf.compareTo(start) >= 0) { stack.push(z); } }
/** * @return the number of items strictly before the current element */ public int headCount(Centroid base) { if (size == 0) { return 0; } else if (left == null) { return leaf.compareTo(base) < 0 ? 1 : 0; } else { if (base.compareTo(leaf) < 0) { return left.headCount(base); } else { return left.size + right.headCount(base); } } }
/** * @return the sum of the size() function for all elements strictly before the current element. */ public long headSum(Centroid base) { if (size == 0) { return 0; } else if (left == null) { return leaf.compareTo(base) < 0 ? count : 0; } else { if (base.compareTo(leaf) <= 0) { return left.headSum(base); } else { return left.count + right.headSum(base); } } }
public void print(int depth) { for (int i = 0; i < depth; i++) { System.out.printf("| "); } int imbalance = Math.abs((left != null ? left.depth : 1) - (right != null ? right.depth : 1)); System.out.printf("%s%s, %d, %d, %d\n", (imbalance > 1 ? "* " : "") + (right != null && leaf.compareTo(right.first()) != 0 ? "+ " : ""), leaf, size, count, this.depth); if (left != null) { left.print(depth + 1); right.print(depth + 1); } } }
/** * @return the smallest element greater than or equal to base. */ public Centroid ceiling(Centroid base) { if (size == 0) { return null; } else if (size == 1) { return base.compareTo(leaf) <= 0 ? leaf : null; } else { if (base.compareTo(leaf) < 0) { Centroid r = left.ceiling(base); if (r == null) { r = right.first(); } return r; } else { return right.ceiling(base); } } }
/** * @return the largest element less than or equal to base */ public Centroid floor(Centroid base) { if (size == 0) { return null; } else { if (size == 1) { return base.compareTo(leaf) >= 0 ? leaf : null; } else { if (base.compareTo(leaf) < 0) { return left.floor(base); } else { Centroid floor = right.floor(base); if (floor == null) { floor = left.last(); } return floor; } } } }
/** * Modify an existing value in the tree subject to the constraint that the change will not alter the * ordering of the tree. * @param x New value to add to Centroid * @param count Weight of new value * @param v The value to modify * @param data The recorded data */ public void move(double x, int count, Centroid v, Iterable<? extends Double> data) { if (size <= 0) { throw new IllegalStateException("Cannot move element of empty tree"); } if (size == 1) { if(leaf != v) { throw new IllegalStateException("Cannot move element that is not in tree"); } leaf.add(x, count, data); } else if (v.compareTo(leaf) < 0) { left.move(x, count, v, data); } else { right.move(x, count, v, data); } this.count += count; }
public void checkBalance() { if (left != null) { if(Math.abs(left.depth() - right.depth()) >= 2) { throw new IllegalStateException("Imbalanced"); } int l = left.depth(); int r = right.depth(); if(depth != Math.max(l, r) + 1){ throw new IllegalStateException( "Depth doesn't match children"); } if(size != left.size + right.size){ throw new IllegalStateException( "Sizes don't match children"); } if(count != left.count + right.count){ throw new IllegalStateException( "Counts don't match children"); } if(leaf.compareTo(right.first()) != 0){ throw new IllegalStateException(String.format( "Split is wrong %.5f != %.5f or %d != %d", leaf.mean(), right.first().mean(), leaf.id(), right.first().id())); } left.checkBalance(); right.checkBalance(); } }
public void add(Centroid centroid) { if (size == 0) { leaf = centroid; depth = 1; count = centroid.count(); size = 1; return; } else if (size == 1) { int order = centroid.compareTo(leaf); if (order < 0) { left = new GroupTree(centroid); right = new GroupTree(leaf); } else if (order > 0) { left = new GroupTree(leaf); right = new GroupTree(centroid); leaf = centroid; } } else if (centroid.compareTo(leaf) < 0) { left.add(centroid); } else { right.add(centroid); } count += centroid.count(); size++; depth = Math.max(left.depth, right.depth) + 1; rebalance(); }