working weight balanced tree (both map & set)

assira.commons/src/main/java/net/ranides/assira/collection/maps/ASortedMap.java

@@ -509,7 +509,7 @@ public abstract class ASortedMap<K, V> extends AMap<K, V> implements SortedMap<K
         }
 
         @Override
-        public final int size() {
+        public int size() {
             return IteratorUtils.size(entryIterator());
         }
 

assira.commons/src/main/java/net/ranides/assira/collection/maps/WBTreeMap.java

@@ -1,7 +1,5 @@
 package net.ranides.assira.collection.maps;
 
-import net.ranides.assira.collection.sets.WBTree;
-import net.ranides.assira.collection.sets.WBTreeSet;
 import net.ranides.assira.generic.CompareUtils;
 
 import java.io.Serializable;
@@ -11,7 +9,7 @@ import java.util.NoSuchElementException;
 
 public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
-    private final WBTree<PEntry<K,V>> tree;
+    private final WBTreeMapBase<K,V> tree;
 
     public WBTreeMap() {
         this(CompareUtils.comparator());
@@ -19,9 +17,9 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
     public WBTreeMap(Comparator<? super K> comparator) {
         super(comparator);
-        this.tree = new WBTree<>((Serializable & Comparator<Entry<K,V>>)(a,b) -> {
+        this.tree = new WBTreeMapBase<>((Serializable & Comparator<? super K>)(a,b) -> {
             try {
-                return comparator.compare(a.getKey(), b.getKey());
+                return comparator.compare(a, b);
             } catch (ClassCastException e) {
                 return -1;
             }
@@ -30,20 +28,17 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
     @Override
     protected Entry<K, V> findEntry(K key) {
-        WBTree.Node<PEntry<K, V>> node = tree.find(new PEntry<>(key, null));
-        return node != null ? node.value : null;
+        return tree.find(key);
     }
 
     @Override
     protected Entry<K, V> firstEntry() {
-        WBTree.Node<PEntry<K, V>> node = tree.first();
-        return node != null ? node.value : null;
+        return tree.first();
     }
 
     @Override
     protected Entry<K, V> lastEntry() {
-        WBTree.Node<PEntry<K, V>> node = tree.last();
-        return node != null ? node.value : null;
+        return tree.last();
     }
 
     @Override
@@ -53,7 +48,9 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
     @Override
     protected ListIterator<Entry<K, V>> entryIterator(K begin) {
-        return new TreeIterator(tree.find(new PEntry<>(begin, null), WBTree.FindMode.EQG));
+        WBTreeMapBase.Node<K, V> after = tree.find(begin, WBTreeMapBase.FindMode.GE);
+        int index = WBTreeMapBase.indexOf(after);
+        return new TreeIterator(index < 0 ? tree.size() : index);
     }
 
     @Override
@@ -68,8 +65,8 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
     @Override
     public boolean remove(Object key, Object value) {
-        WBTree.Node<PEntry<K, V>> prev = tree.find(new PEntry<>((K) key, (V) value));
-        if(prev != null && CompareUtils.equals(prev.value.getValue(), value)) {
+        WBTreeMapBase.Node<K, V> prev = tree.find((K)key);
+        if(prev != null && CompareUtils.equals(prev.value, value)) {
             tree.removeNode(prev);
             return true;
         }
@@ -78,14 +75,19 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
     @Override
     public V put(K key, V value) {
-        WBTree.Node<PEntry<K, V>> prev = tree.add(new PEntry<>(key, value));
-        return (prev != null) ? prev.value.getValue() : null;
+        WBTreeMapBase.Node<K, V> prev = tree.find(key);
+        if(prev != null) {
+            return prev.setValue(value);
+        } else {
+            tree.add(key, value);
+            return null;
+        }
     }
 
     @Override
     public V remove(Object key) {
-        WBTree.Node<PEntry<K, V>> prev = tree.remove(new PEntry<K, V>((K) key, null));
-        return (prev != null) ? prev.value.getValue() : null;
+        WBTreeMapBase.Node<K, V> prev = tree.remove((K) key);
+        return (prev != null) ? prev.value : null;
     }
 
     @Override
@@ -99,6 +101,33 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
             super(begin, bottom, end, top);
         }
 
+        @Override
+        public int size() {
+            WBTreeMapBase.Node<K,V> first = $first();
+            WBTreeMapBase.Node<K,V> last = $last();
+            if(first == null || last == null) {
+                return 0;
+            } else {
+                return WBTreeMapBase.indexOf(last) - WBTreeMapBase.indexOf(first) + 1;
+            }
+        }
+
+        private WBTreeMapBase.Node<K,V> $first() {
+            WBTreeMapBase.Node<K,V> out = bottom ? tree.first() : tree.find(begin, WBTreeMapBase.FindMode.GE);
+            if(out == null || !inside(out.key)) {
+                return null;
+            }
+            return out;
+        }
+
+        public WBTreeMapBase.Node<K,V> $last() {
+            WBTreeMapBase.Node<K,V> out = top ? tree.last() : tree.find(end, WBTreeMapBase.FindMode.LT);
+            if(out == null || !inside(out.key)) {
+                return null;
+            }
+            return out;
+        }
+
         @Override
         protected ListIterator<Entry<K,V>> entryIterator() {
             return new SIterator();
@@ -106,7 +135,20 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
         @Override
         protected ListIterator<Entry<K,V>> entryIterator(K from) {
-            return new SIterator(tree.find(new PEntry<>(WBTreeMap.this.higher(from, begin), null)));
+            WBTreeMapBase.Node<K,V> fromNode;
+            if(bottom) {
+                fromNode = tree.find(from, WBTreeMapBase.FindMode.GT);
+            }
+            else if(compare(from, begin) > 0) {
+                fromNode = tree.find(from, WBTreeMapBase.FindMode.GT);
+            } else {
+                fromNode = tree.find(begin, WBTreeMapBase.FindMode.GE);
+            }
+            if(fromNode == null) {
+                return new SIterator(size());
+            } else {
+                return new SIterator(WBTreeMapBase.indexOf(fromNode) - WBTreeMapBase.indexOf($first()));
+            }
         }
 
         @Override
@@ -119,169 +161,114 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
         @Override
         protected Entry<K,V> firstEntry() {
-            // If this submap goes end -infinity, we return the main map first entry; otherwise, we locate the start of the map.
-            WBTree.Node<PEntry<K, V>> entry;
-            if (bottom) {
-                entry = tree.first();
-            } else {
-                entry = tree.find(new PEntry<>(begin, null));
-                // If we find either the start or something greater we're OK.
-                if (entry!=null && compare(entry.value.getKey(), begin) < 0) {
-                    entry = WBTree.next(entry);
-                }
-            }
-            // Finally, if this subset doesn't go end infinity, we check that the resulting key isn't greater than the end.
-            if (entry == null || !top && compare(entry.value.getKey(), end) >= 0) {
-                return null;
+            WBTreeMapBase.Node<K,V> out = $first();
+            if(out == null) {
+                throw new NoSuchElementException();
             }
-            return entry.value;
+            return out;
         }
 
         @Override
-        protected Entry lastEntry() {
-            if (tree == null) {
-                return null;
-            }
-            // If this submap goes end infinity, we return the main map last entry; otherwise, we locate the end of the map.
-            WBTree.Node<PEntry<K, V>> entry;
-            if (top) {
-                entry = tree.last();
-            } else {
-                entry = tree.find(new PEntry<>(begin, null));
-                // If we find something smaller than the end we're OK.
-                if (entry!=null && compare(entry.value.getKey(), end) >= 0) {
-                    entry = WBTree.prev(entry);
-                }
-            }
-            // Finally, if this subset doesn't go end -infinity, we check that the resulting key isn't smaller than the start.
-            if (entry == null || !bottom && compare(entry.value.getKey(), begin) < 0) {
-                return null;
+        protected Entry<K,V> lastEntry() {
+            WBTreeMapBase.Node<K,V> out = $last();
+            if(out == null) {
+                throw new NoSuchElementException();
             }
-            return entry.value;
+            return out;
         }
 
+
+
         private class SIterator implements ListIterator<Entry<K,V>> {
 
-            private WBTree.Node<PEntry<K,V>> current;
-            private WBTree.Node<PEntry<K,V>> prev;
-            private WBTree.Node<PEntry<K,V>> next;
-            private int offset;
+            private int index;
+            private int last;
+
+            private final int minIndex;
+            private int maxIndex;
+
+            private final int shift;
 
             public SIterator() {
-                if(bottom) {
-                    this.current = tree.first();
-                    this.next = current;
-                    this.offset = 0;
+                this(0);
+            }
+
+            public SIterator(int index) {
+                this.index = index;
+                this.last = -1;
+                this.shift = index;
+
+                WBTreeMapBase.Node<K,V> first = $first();
+                WBTreeMapBase.Node<K,V> last = $last();
+                if(first == null || last == null) {
+                    this.minIndex = 0;
+                    this.maxIndex = 0;
                 } else {
-                    this.current = tree.find(new PEntry<>(begin,null), WBTree.FindMode.LE);
-                    this.next = null;
-                    this.offset = WBTree.indexOf(current);
+                    this.minIndex = WBTreeMapBase.indexOf(first);
+                    this.maxIndex = WBTreeMapBase.indexOf(last) + 1;
                 }
             }
 
-            public SIterator(WBTree.Node<PEntry<K,V>> after) {
-                this.current = after;
-                this.next = null;
+            @Override
+            public final boolean hasNext() {
+                return minIndex+index < maxIndex;
             }
 
             @Override
-            public boolean hasNext() {
-                if(next != null) {
-                    return true;
-                }
-                if(current == null) {
-                    return false;
-                }
-                next = WBTree.next(current);
-
-                return next != null && inside(next.value.getKey());
+            public final boolean hasPrevious() {
+                return minIndex+index > minIndex;
             }
 
             @Override
             public Entry<K,V> next() {
-                if(next == null) {
-                    if(current == null) {
-                        throw new NoSuchElementException();
-                    }
-                    next = WBTree.next(current);
-                }
-
-                prev = current;
-                current = next;
-                next = null;
-
-                if(current == null || !inside(current.value.getKey())) {
+                if (!hasNext()) {
                     throw new NoSuchElementException();
                 }
-
-                return current.value;
-            }
-
-            @Override
-            public boolean hasPrevious() {
-                if(prev != null) {
-                    return true;
-                }
-                if(current == null) {
-                    return false;
-                }
-                prev = WBTree.prev(current);
-
-                return prev != null && inside(prev.value.getKey());
+                return tree.at((last = index++) + minIndex);
             }
 
             @Override
             public Entry<K,V> previous() {
-                if(prev == null) {
-                    if(current == null) {
-                        throw new NoSuchElementException();
-                    }
-                    prev = WBTree.prev(current);
-                }
-
-                next = current;
-                current = prev;
-                prev = null;
-
-                if(current == null || !inside(current.value.getKey())) {
+                if (!hasPrevious()) {
                     throw new NoSuchElementException();
                 }
-
-                return current.value;
+                return tree.at((last = --index) + minIndex);
             }
 
             @Override
-            public int nextIndex() {
-                // @todo optimize, we should store index
-                return offset + WBTree.indexOf(current);
+            public final int nextIndex() {
+                return index - minIndex - shift;
             }
 
             @Override
-            public int previousIndex() {
-                // @todo optimize, we should store index
-                return offset + WBTree.indexOf(current) - 1;
+            public final int previousIndex() {
+                return index - minIndex - shift - 1;
             }
 
             @Override
-            public void remove() {
-                if(current == null) {
-                    throw new IllegalStateException();
-                }
-                prev = WBTree.prev(current);
-                next = WBTree.next(current);
-
-                tree.removeNode(current);
-                current = null;
+            public final void add(Entry<K,V> value) {
+                throw new UnsupportedOperationException();
             }
 
             @Override
-            public void set(Entry<K,V> k) {
+            public final void set(Entry<K,V> value) {
                 throw new UnsupportedOperationException();
             }
 
             @Override
-            public void add(Entry<K,V> k) {
-                throw new UnsupportedOperationException();
+            public final void remove() {
+                if (last == -1) {
+                    throw new IllegalStateException();
+                }
+                tree.removeNode(tree.at(last + minIndex));
+                // If the last operation was a next(),
+                // we are removing an element *before* us,
+                // and we must decrease pos correspondingly.
+                if (last < index) {
+                    index--;
+                }
+                last = -1;
+                maxIndex--;
             }
         }
 
@@ -289,148 +276,82 @@ public class WBTreeMap<K,V> extends ASortedMap<K,V> {
 
     private class TreeIterator implements ListIterator<Entry<K,V>> {
 
-        private WBTree.Node<PEntry<K,V>> current;
-        private WBTree.Node<PEntry<K,V>> prev;
-        private WBTree.Node<PEntry<K,V>> next;
+        private int index;
+
+        private int last;
+
+        private final int offset;
 
         public TreeIterator() {
-            this.current = tree.first();
-            this.next = current;
+            this(0);
         }
 
-        public TreeIterator(WBTree.Node<PEntry<K,V>> after) {
-            this.current = after;
-            this.next = null;
+        public TreeIterator(int index) {
+            this.last = -1;
+            this.index = index;
+            this.offset = index;
         }
 
         @Override
-        public boolean hasNext() {
-            if(next != null) {
-                return true;
-            }
-            if(current == null) {
-                return false;
-            }
-            next = WBTree.next(current);
-
-            return next != null;
+        public final boolean hasNext() {
+            return index < size();
         }
 
         @Override
-        public Entry<K,V> next() {
-            if(next == null) {
-                if(current == null) {
-                    throw new NoSuchElementException();
-                }
-                next = WBTree.next(current);
-            }
-
-            prev = current;
-            current = next;
-            next = null;
-
-            if(current == null) {
-                throw new NoSuchElementException();
-            }
-
-            return current.value;
+        public final boolean hasPrevious() {
+            return index > 0;
         }
 
         @Override
-        public boolean hasPrevious() {
-            if(prev != null) {
-                return true;
-            }
-            if(current == null) {
-                return false;
+        public final Entry<K,V> next() {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
             }
-            prev = WBTree.prev(current);
-
-            return prev != null;
+            return tree.at(last = index++);
         }
 
         @Override
-        public Entry<K,V> previous() {
-            if(prev == null) {
-                if(current == null) {
-                    throw new NoSuchElementException();
-                }
-                prev = WBTree.prev(current);
-            }
-
-            next = current;
-            current = prev;
-            prev = null;
-
-            if(current == null) {
+        public final Entry<K,V> previous() {
+            if (!hasPrevious()) {
                 throw new NoSuchElementException();
             }
-
-            return current.value;
-        }
-
-        @Override
-        public int nextIndex() {
-            // @todo optimize, we should store index
-            return WBTree.indexOf(current);
+            return tree.at(last = --index);
         }
 
         @Override
-        public int previousIndex() {
-            // @todo optimize, we should store index
-            return WBTree.indexOf(current) - 1;
+        public final int nextIndex() {
+            return index - offset;
         }
 
         @Override
-        public void remove() {
-            if(current == null) {
-                throw new IllegalStateException();
-            }
-            prev = WBTree.prev(current);
-            next = WBTree.next(current);
-
-            tree.removeNode(current);
-            current = null;
+        public final int previousIndex() {
+            return index - offset - 1;
         }
 
         @Override
-        public void set(Entry<K,V> k) {
+        public final void add(Entry<K,V> value) {
             throw new UnsupportedOperationException();
         }
 
         @Override
-        public void add(Entry<K,V> k) {
+        public final void set(Entry<K,V> value) {
             throw new UnsupportedOperationException();
         }
-    }
-
-    protected static final class PEntry<K, V> extends AEntry<K, V> implements Serializable {
-
-        private static final long serialVersionUID = 2L;
-
-        private final K key;
-        private V value;
-
-        public PEntry(K key, V value) {
-            this.key = key;
-            this.value = value;
-        }
 
         @Override
-        public K getKey() {
-            return key;
-        }
-
-        @Override
-        public V getValue() {
-            return value;
-        }
-
-        @Override
-        public V setValue(V value) {
-            V prev = this.value;
-            this.value = value;
-            return prev;
+        public final void remove() {
+            if (last == -1) {
+                throw new IllegalStateException();
+            }
+            tree.removeNode(tree.at(last));
+            // If the last operation was a next(),
+            // we are removing an element *before* us,
+            // and we must decrease pos correspondingly.
+            if (last < index) {
+                index--;
+            }
+            last = -1;
         }
     }
+
 }

assira.commons/src/main/java/net/ranides/assira/collection/maps/WBTreeMapBase.java

@@ -0,0 +1,427 @@
+/**
+ * This work is based on https://deegen1.github.io/wbtree/index.html
+ *
+ * WBTreeSetBase.py - v1.01
+ * Copyright 2022 Alec Dee - MIT license - SPDX: MIT
+ * deegen1.github.io - akdee144@gmail.com
+ *
+ * This class implements reasonable interface for weight balanced tree, but it is not java-opinionated
+ * For example, it does not care about Iterators, sub-sets, submaps, spliterators etc
+ *
+ * We use this structure inside SortedSet and SortedMap classes.
+ * WBTreeSet and WBTreeMap implement all java idioms, when core algorithmic stuff is stored here.
+ *
+ * Such decomposition makes few things somewhat less effective (e.g. map contains double-wrapped entries) but
+ * it allows us to synchronize with python original source code.
+ */
+package net.ranides.assira.collection.maps;
+
+import net.ranides.assira.generic.CompareUtils;
+
+import java.io.Serializable;
+import java.util.Comparator;
+import java.util.function.Consumer;
+
+public class WBTreeMapBase<K,V> implements Serializable {
+
+    private static final long serialVersionUID = 2L;
+
+    public enum FindMode {
+        EQ,
+        EQG,
+        LT,
+        LE,
+        GT,
+        GE
+    }
+
+    private final Comparator<? super K> cmp;
+    private final boolean unique;
+    private Node<K,V> root;
+
+    public WBTreeMapBase() {
+        this(CompareUtils.comparator(), true);
+    }
+
+    public WBTreeMapBase(Comparator<? super K> cmp, boolean unique) {
+        this.cmp = cmp;
+        this.unique = unique;
+        this.root = null;
+    }
+
+    public void clear() {
+        this.root = null;
+    }
+
+    public int size() {
+        return weight(root);
+    }
+
+    public Node<K,V> at(int i) {
+        Node<K,V> node = root;
+        int weight = size();
+        if(i<0) {
+            i += weight;
+        }
+        if(i<0 || i>=weight) {
+            return null;
+        }
+        while (true) {
+            Node<K,V> left = node.left;
+            int lw = weight(left);
+            if(i > lw) {
+                i -= lw +1;
+                node = node.right;
+            }
+            else if(i == lw) {
+                break;
+            } else {
+                node = left;
+            }
+        }
+        return node;
+    }
+
+    public void forEach(Consumer<K> consumer) {
+        forEachNode(node -> consumer.accept(node.key));
+    }
+
+    void forEachNode(Consumer<Node<K,V>> consumer) {
+        Node<K,V> node = first();
+        while(node != null) {
+            consumer.accept(node);
+            node = next(node);
+        }
+    }
+
+    public Node<K,V> first() {
+        Node<K,V> node = root;
+        Node<K,V> ret = null;
+
+        while(node != null) {
+            ret = node;
+            node = node.left;
+        }
+
+        return ret;
+    }
+
+    public Node<K,V> last() {
+        Node<K,V> node = root;
+        Node<K,V> ret = null;
+
+        while (node != null) {
+            ret = node;
+            node = node.right;
+        }
+
+        return ret;
+    }
+
+    public Node<K,V> find(K value) {
+        return find(value, FindMode.EQ);
+    }
+
+    public Node<K,V> find(K value, FindMode mode) {
+        Node<K,V> node = root;
+        Node<K,V> ret = null;
+
+        boolean lset = mode == FindMode.LT || mode == FindMode.LE;
+        boolean rset = mode == FindMode.GT || mode == FindMode.GE;
+        boolean eset = mode == FindMode.EQ || mode == FindMode.EQG || mode == FindMode.LE || mode == FindMode.GE;
+        boolean eright = mode == FindMode.EQG || mode == FindMode.LE || mode == FindMode.GT;
+
+        while(node != null) {
+            int c = cmp.compare(node.key, value);
+            if(c < 0) {
+                if(lset) {
+                    ret = node;
+                }
+                node = node.right;
+            }
+            else if(c > 0) {
+                if(rset) {
+                    ret = node;
+                }
+                node = node.left;
+            }
+            else {
+                if(eset) {
+                    ret = node;
+                    if(unique) {
+                        break;
+                    }
+                }
+                node = eright ? node.right : node.left;
+            }
+        }
+
+        return ret;
+    }
+
+    public Node<K,V> add(K key, V value) {
+        Node<K,V> prev = null;
+        Node<K,V> node = root;
+        int c = 0;
+        while(node != null) {
+            c = cmp.compare(node.key, key);
+            if(c == 0 && this.unique) {
+                node.key = key;
+                node.value = value;
+                return node;
+            }
+            prev = node;
+            if(c <= 0) {
+                node = node.right;
+            } else {
+                node = node.left;
+            }
+        }
+        node = new Node<>(key, value);
+        node.parent = prev;
+        if(prev == null) {
+            this.root = node;
+        } else {
+            if(c <=0) {
+                prev.right = node;
+            } else {
+                prev.left = node;
+            }
+            rebalance(prev);
+        }
+        return node;
+    }
+
+    public Node<K,V> remove(K value) {
+        Node<K,V> node = find(value);
+        if(node != null) {
+            removeNode(node);
+        }
+        return node;
+    }
+
+    public void removeNode(Node<K,V> node) {
+        Node<K,V> p = node.parent;
+        Node<K,V> l = node.left;
+        Node<K,V> r = node.right;
+        Node<K,V> bal = null;
+        Node<K,V> next = null;
+
+        if(r == null) {
+            // case 1
+            bal = p;
+            next = l;
+            l = null;
+        }
+        else if(r.left != null) {
+            // case 3
+            next = r;
+            while(next.left != null) {
+                bal = next;
+                next = next.left;
+            }
+            Node<K,V> c = next.right;
+            bal.left = c;
+            if(c != null) {
+                c.parent = bal;
+            }
+        }
+        else {
+            // case 2
+            bal = r;
+            next = r;
+            r = null;
+        }
+        // replace node with next
+        if(p == null) {
+            root = next;
+        }
+        else if(p.left == node) {
+            p.left = next;
+        }
+        else {
+            p.right = next;
+        }
+
+        if(next != null) {
+            next.parent = p;
+            if(l != null) {
+                next.left = l;
+                l.parent = next;
+            }
+            if(r != null) {
+                next.right = r;
+                r.parent = next;
+            }
+        }
+
+        rebalance(bal);
+    }
+
+    void rebalance(Node<K,V> next) {
+        while (next != null) {
+            Node<K,V> node = next;
+            Node<K,V> orig = next;
+            next = node.parent;
+
+            Node<K,V> l = node.left;
+            Node<K,V> r = node.right;
+            int lw = weight(l);
+            int rw = weight(r);
+
+            if(rw * 5 + 2< lw * 2) {
+                if(weight(l.left) * 5 < lw *2) {
+                    node.left = rotateLeft(l);
+                }
+                node = rotateRight(node);
+            }
+            else if (lw * 5 + 2 < rw * 2) {
+                if( weight(r.right) * 5 < rw*2) {
+                    node.right = rotateRight(r);
+                }
+                node = rotateLeft(node);
+            }
+            else {
+                node.weight = lw + rw + 1;
+                continue;
+            }
+            if(next == null) {
+                this.root = node;
+            }
+            else if(next.left == orig) {
+                next.left = node;
+            }
+            else {
+                next.right = node;
+            }
+        }
+    }
+
+    public static int weight(Node<?,?> node) {
+        return node != null ? node.weight : 0;
+    }
+
+    public static <K,V> Node<K,V> next(Node<K,V> node) {
+        Node<K,V> child = node.right;
+        if(child != null) {
+            while(child != null) {
+                node = child;
+                child = child.left;
+            }
+        } else {
+            while(node != null && node.right == child) {
+                child = node;
+                node = node.parent;
+            }
+        }
+        return node;
+    }
+
+    public static <K,V> Node<K,V> prev(Node<K,V> node) {
+        Node<K,V> child = node.left;
+        if(child != null) {
+            while(child != null) {
+                node = child;
+                child = child.right;
+            }
+        } else {
+            while(node != null && node.left == child) {
+                child = node;
+                node = node.parent;
+            }
+        }
+        return node;
+    }
+
+    public static <K,V> Node<K,V> rotateLeft(Node<K,V> a) {
+        Node<K,V> b=a.right;
+        Node<K,V> r=b.left;
+        b.parent=a.parent;
+        b.left=a;
+        a.parent=b;
+        a.right=r;
+        if (r != null) {
+            r.parent = a;
+        }
+        recomputeWeight(a);
+        recomputeWeight(b);
+        return b;
+    }
+
+    public static <K,V> Node<K,V> rotateRight(Node<K,V> a) {
+        Node<K,V> b = a.left;
+        Node<K,V> l = b.right;
+        b.parent = a.parent;
+        b.right = a;
+        a.parent = b;
+        a.left = l;
+        if (l != null) {
+            l.parent = a;
+        }
+        recomputeWeight(a);
+        recomputeWeight(b);
+        return b;
+    }
+
+    public static void recomputeWeight(Node<?,?> node) {
+        node.weight = 1 + weight(node.left) + weight(node.right);
+    }
+
+    public static int indexOf(Node<?,?> node) {
+        if(node == null) {
+            return -1;
+        }
+        int i = -1;
+        Node<?,?> prev = node.right;
+        while (node != null) {
+            if (node.right == prev) {
+                i += weight(node.left) + 1;
+            }
+            prev = node;
+            node = node.parent;
+        }
+        return i;
+    }
+
+    public static class Node<K, V> extends AMap.AEntry<K, V> implements Serializable {
+
+        private static final long serialVersionUID = 2L;
+
+        public int weight;
+        public Node<K,V> parent;
+        public Node<K,V> left;
+        public Node<K,V> right;
+        public K key;
+        public V value;
+
+        public Node(K key, V value) {
+            this.weight = 1;
+            this.parent = null;
+            this.left = null;
+            this.right = null;
+            this.key = key;
+            this.value = value;
+        }
+
+        @Override
+        public K getKey() {
+            return key;
+        }
+
+        @Override
+        public V getValue() {
+            return value;
+        }
+
+        @Override
+        public V setValue(V value) {
+            V prev = this.value;
+            this.value = value;
+            return prev;
+        }
+    }
+
+
+
+
+}

assira.commons/src/main/java/net/ranides/assira/collection/sets/WBTreeSet.java

@@ -8,7 +8,7 @@ import java.util.NoSuchElementException;
 
 public class WBTreeSet<K> extends ASortedSet<K> {
 
-    private final WBTree<K> tree;
+    private final WBTreeSetBase<K> tree;
 
     public WBTreeSet() {
         this(CompareUtils.comparator());
@@ -16,7 +16,7 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
     public WBTreeSet(Comparator<? super K> comparator) {
         super(comparator);
-        this.tree = new WBTree<>(comparator, true);
+        this.tree = new WBTreeSetBase<>(comparator, true);
     }
 
     @Override
@@ -48,8 +48,8 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
     @Override
     public ListIterator<K> iterator(K begin) {
-        WBTree.Node<K> after = tree.find(begin, WBTree.FindMode.GT);
-        int index = WBTree.indexOf(after);
+        WBTreeSetBase.Node<K> after = tree.find(begin, WBTreeSetBase.FindMode.GT);
+        int index = WBTreeSetBase.indexOf(after);
         return new TreeIterator(index < 0 ? tree.size() : index);
     }
 
@@ -65,7 +65,7 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
     @Override
     public K first() {
-        WBTree.Node<K> node = tree.first();
+        WBTreeSetBase.Node<K> node = tree.first();
         if(node == null) {
             throw new NoSuchElementException();
         }
@@ -74,7 +74,7 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
     @Override
     public K last() {
-        WBTree.Node<K> node = tree.last();
+        WBTreeSetBase.Node<K> node = tree.last();
         if(node == null) {
             throw new NoSuchElementException();
         }
@@ -88,13 +88,13 @@ public class WBTreeSet<K> extends ASortedSet<K> {
         }
 
         @Override
-        public final int size() {
-            WBTree.Node<K> first = $first();
-            WBTree.Node<K> last = $last();
+        public int size() {
+            WBTreeSetBase.Node<K> first = $first();
+            WBTreeSetBase.Node<K> last = $last();
             if(first == null || last == null) {
                 return 0;
             } else {
-                return WBTree.indexOf(last) - WBTree.indexOf(first) + 1;
+                return WBTreeSetBase.indexOf(last) - WBTreeSetBase.indexOf(first) + 1;
             }
         }
 
@@ -105,32 +105,32 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
         @Override
         public ListIterator<K> iterator(K from) {
-            WBTree.Node<K> fromNode;
+            WBTreeSetBase.Node<K> fromNode;
             if(bottom) {
-                fromNode = tree.find(from, WBTree.FindMode.GT);
+                fromNode = tree.find(from, WBTreeSetBase.FindMode.GT);
             }
             else if(compare(from, begin) > 0) {
-                fromNode = tree.find(from, WBTree.FindMode.GT);
+                fromNode = tree.find(from, WBTreeSetBase.FindMode.GT);
             } else {
-                fromNode = tree.find(begin, WBTree.FindMode.GE);
+                fromNode = tree.find(begin, WBTreeSetBase.FindMode.GE);
             }
             if(fromNode == null) {
                 return new SIterator(size());
             } else {
-                return new SIterator(WBTree.indexOf(fromNode) - WBTree.indexOf($first()));
+                return new SIterator(WBTreeSetBase.indexOf(fromNode) - WBTreeSetBase.indexOf($first()));
             }
         }
 
-        private WBTree.Node<K> $first() {
-            WBTree.Node<K> out = bottom ? tree.first() : tree.find(begin, WBTree.FindMode.GE);
+        private WBTreeSetBase.Node<K> $first() {
+            WBTreeSetBase.Node<K> out = bottom ? tree.first() : tree.find(begin, WBTreeSetBase.FindMode.GE);
             if(out == null || !inside(out.value)) {
                 return null;
             }
             return out;
         }
 
-        public WBTree.Node<K> $last() {
-            WBTree.Node<K> out = top ? tree.last() : tree.find(end, WBTree.FindMode.LT);
+        public WBTreeSetBase.Node<K> $last() {
+            WBTreeSetBase.Node<K> out = top ? tree.last() : tree.find(end, WBTreeSetBase.FindMode.LT);
             if(out == null || !inside(out.value)) {
                 return null;
             }
@@ -139,7 +139,7 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
         @Override
         public K first() {
-            WBTree.Node<K> out = $first();
+            WBTreeSetBase.Node<K> out = $first();
             if(out == null) {
                 throw new NoSuchElementException();
             }
@@ -148,7 +148,7 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
         @Override
         public K last() {
-            WBTree.Node<K> out = $last();
+            WBTreeSetBase.Node<K> out = $last();
             if(out == null) {
                 throw new NoSuchElementException();
             }
@@ -174,14 +174,14 @@ public class WBTreeSet<K> extends ASortedSet<K> {
                 this.last = -1;
                 this.shift = index;
 
-                WBTree.Node<K> first = $first();
-                WBTree.Node<K> last = $last();
+                WBTreeSetBase.Node<K> first = $first();
+                WBTreeSetBase.Node<K> last = $last();
                 if(first == null || last == null) {
                     this.minIndex = 0;
                     this.maxIndex = 0;
                 } else {
-                    this.minIndex = WBTree.indexOf(first);
-                    this.maxIndex = WBTree.indexOf(last) + 1;
+                    this.minIndex = WBTreeSetBase.indexOf(first);
+                    this.maxIndex = WBTreeSetBase.indexOf(last) + 1;
                 }
             }
 

assira.commons/src/main/java/net/ranides/assira/collection/sets/WBTree.java

@@ -1,7 +1,7 @@
 /**
  * This work is based on https://deegen1.github.io/wbtree/index.html
  *
- * WBTree.py - v1.01
+ * WBTreeSetBase.py - v1.01
  * Copyright 2022 Alec Dee - MIT license - SPDX: MIT
  * deegen1.github.io - akdee144@gmail.com
  *
@@ -22,7 +22,7 @@ import java.io.Serializable;
 import java.util.Comparator;
 import java.util.function.Consumer;
 
-public class WBTree<V> implements Serializable {
+public class WBTreeSetBase<V> implements Serializable {
 
     private static final long serialVersionUID = 2L;
 
@@ -39,11 +39,11 @@ public class WBTree<V> implements Serializable {
     private final boolean unique;
     private Node<V> root;
 
-    public WBTree() {
+    public WBTreeSetBase() {
         this(CompareUtils.comparator(), true);
     }
 
-    public WBTree(Comparator<? super V> cmp, boolean unique) {
+    public WBTreeSetBase(Comparator<? super V> cmp, boolean unique) {
         this.cmp = cmp;
         this.unique = unique;
         this.root = null;

assira.commons/src/test/java/net/ranides/assira/collection/maps/WBTreeMapTest.java

@@ -49,7 +49,7 @@ public class WBTreeMapTest {
         tree.forEach((k,v) -> values2.add(k));
         assertEquals(Arrays.asList(1,2,4,5), values2);
 
-        assertEquals(4, tree.entryIterator(2).next().getKey().intValue());
+        assertEquals(2, tree.entryIterator(2).next().getKey().intValue());
 
         assertArrayEquals(new Integer[]{1,2,4,5}, tree.keySet().toArray(new Integer[0]));
     }

assira.commons/src/test/java/net/ranides/assira/collection/sets/WBTreeTest.java

@@ -8,11 +8,11 @@ import java.util.List;
 
 import static org.junit.Assert.*;
 
-public class WBTreeTest {
+public class WBTreeSetBaseTest {
 
     @Test
     public void basic() {
-        WBTree<Integer> tree = new WBTree<>();
+        WBTreeSetBase<Integer> tree = new WBTreeSetBase<>();
 
         tree.add(5);
         tree.add(3);