Initial commit

tree_adjust.go

@@ -0,0 +1,74 @@
+package stl
+
+type treeConnectType int8
+
+const (
+	treeLeft treeConnectType = iota
+	treeRight
+)
+
+func (t treeConnectType) invert() treeConnectType {
+	return 1 - t
+}
+
+func (son *Node[K, V, C]) tell() treeConnectType {
+	if son.parent == nil {
+		return treeLeft
+	}
+
+	if son.parent.lson == son {
+		return treeLeft
+	} else {
+		return treeRight
+	}
+}
+
+func (n *Node[K, V, C]) get(t treeConnectType) *Node[K, V, C] {
+	if t == treeLeft {
+		return n.lson
+	} else {
+		return n.rson
+	}
+}
+
+func (parent *Node[K, V, C]) connect(son *Node[K, V, C], t treeConnectType) {
+	if son != nil {
+		son.parent = parent
+	}
+	if parent != nil {
+		if t == treeLeft {
+			parent.lson = son
+		} else {
+			parent.rson = son
+		}
+	}
+}
+
+// rotate the node up
+func (node *Node[K, V, C]) rotate(root **Node[K, V, C]) {
+	if node.parent == nil {
+		// is root, sanity check
+		return
+	}
+
+	t := node.tell()
+
+	f := node.parent
+	b := node.get(t.invert())
+
+	f.parent.connect(node, f.tell())
+	node.connect(f, t.invert())
+	f.connect(b, t)
+
+	if node.parent == nil {
+		// new root
+		*root = node
+	}
+}
+
+// adjust the new node in the tree Treap style.
+func (node *Node[K, V, C]) adjust(root **Node[K, V, C]) {
+	for node.parent != nil && node.parent.bal > node.bal {
+		node.rotate(root)
+	}
+}