diff --git a/priority_queue/vbt.go b/priority_queue/vbt.go index fab0650..75dcceb 100644 --- a/priority_queue/vbt.go +++ b/priority_queue/vbt.go @@ -864,6 +864,20 @@ func (tree *vbTree) fixSizeWithRemove(cur *Node) { if rs >= ls*2+factor || ls >= rs*2+factor { cur = tree.fixSize(cur, ls, rs) } + } else if cur.size == 3 { + if cur.children[0] == nil { + if cur.children[1].children[0] == nil { + cur = tree.lrotate3(cur) + } else { + cur = tree.lrrotate3(cur) + } + } else if cur.children[1] == nil { + if cur.children[0].children[1] == nil { + cur = tree.rrotate3(cur) + } else { + cur = tree.rlrotate3(cur) + } + } } cur = cur.parent } diff --git a/priority_queuekey/vbt.go b/priority_queuekey/vbt.go index a5d45f6..c627b05 100644 --- a/priority_queuekey/vbt.go +++ b/priority_queuekey/vbt.go @@ -867,15 +867,15 @@ func (tree *vbTree) fixSizeWithRemove(cur *Node) { } else if cur.size == 3 { if cur.children[0] == nil { if cur.children[1].children[0] == nil { - tree.lrotate3(cur) + cur = tree.lrotate3(cur) } else { - tree.lrrotate3(cur) + cur = tree.lrrotate3(cur) } } else if cur.children[1] == nil { if cur.children[0].children[1] == nil { - tree.rrotate3(cur) + cur = tree.rrotate3(cur) } else { - tree.rlrotate3(cur) + cur = tree.rlrotate3(cur) } } }