75 lines
1.2 KiB
Go
75 lines
1.2 KiB
Go
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)
|
|
}
|
|
}
|