- avl tree: expose children and parent nodes, can be useful when extending structure to have access to these

pull/51/head
Emir Pasic 9 years ago
parent f480e9419a
commit 59734f8164

@ -19,18 +19,18 @@ func assertTreeImplementation() {
// Tree holds elements of the AVL tree. // Tree holds elements of the AVL tree.
type Tree struct { type Tree struct {
Root *Node Root *Node // Root node
size int Comparator utils.Comparator // Key comparator
Comparator utils.Comparator size int // Total number of keys in the tree
} }
// Node is a single element within the tree // Node is a single element within the tree
type Node struct { type Node struct {
Key interface{} Key interface{}
Value interface{} Value interface{}
c [2]*Node Parent *Node // Parent node
p *Node Children [2]*Node // Children nodes
b int8 b int8
} }
// NewWith instantiates an AVL tree with the custom comparator. // NewWith instantiates an AVL tree with the custom comparator.
@ -80,9 +80,9 @@ func (t *Tree) Get(key interface{}) (value interface{}, found bool) {
case cmp == 0: case cmp == 0:
return n.Value, true return n.Value, true
case cmp < 0: case cmp < 0:
n = n.c[0] n = n.Children[0]
case cmp > 0: case cmp > 0:
n = n.c[1] n = n.Children[1]
} }
} }
return nil, false return nil, false
@ -105,10 +105,10 @@ func (t *Tree) Floor(key interface{}) (floor *Node, found bool) {
case c == 0: case c == 0:
return n, true return n, true
case c < 0: case c < 0:
n = n.c[0] n = n.Children[0]
case c > 0: case c > 0:
floor, found = n, true floor, found = n, true
n = n.c[1] n = n.Children[1]
} }
} }
if found { if found {
@ -135,9 +135,9 @@ func (t *Tree) Ceiling(key interface{}) (floor *Node, found bool) {
return n, true return n, true
case c < 0: case c < 0:
floor, found = n, true floor, found = n, true
n = n.c[0] n = n.Children[0]
case c > 0: case c > 0:
n = n.c[1] n = n.Children[1]
} }
} }
if found { if found {
@ -154,7 +154,7 @@ func (t *Tree) Put(key interface{}, value interface{}) {
q := *qp q := *qp
if q == nil { if q == nil {
t.size++ t.size++
*qp = &Node{Key: key, Value: value, p: p} *qp = &Node{Key: key, Value: value, Parent: p}
return true return true
} }
@ -172,7 +172,7 @@ func (t *Tree) Put(key interface{}, value interface{}) {
} }
a := (c + 1) / 2 a := (c + 1) / 2
var fix bool var fix bool
fix = put(q, &q.c[a]) fix = put(q, &q.Children[a])
if fix { if fix {
return putFix(int8(c), qp) return putFix(int8(c), qp)
} }
@ -195,14 +195,14 @@ func (t *Tree) Remove(key interface{}) {
c := t.Comparator(key, q.Key) c := t.Comparator(key, q.Key)
if c == 0 { if c == 0 {
t.size-- t.size--
if q.c[1] == nil { if q.Children[1] == nil {
if q.c[0] != nil { if q.Children[0] != nil {
q.c[0].p = q.p q.Children[0].Parent = q.Parent
} }
*qp = q.c[0] *qp = q.Children[0]
return true return true
} }
fix := removeMin(&q.c[1], &q.Key, &q.Value) fix := removeMin(&q.Children[1], &q.Key, &q.Value)
if fix { if fix {
return removeFix(-1, qp) return removeFix(-1, qp)
} }
@ -215,7 +215,7 @@ func (t *Tree) Remove(key interface{}) {
c = 1 c = 1
} }
a := (c + 1) / 2 a := (c + 1) / 2
fix := remove(&q.c[a]) fix := remove(&q.Children[a])
if fix { if fix {
return removeFix(int8(-c), qp) return removeFix(int8(-c), qp)
} }
@ -227,16 +227,16 @@ func (t *Tree) Remove(key interface{}) {
func removeMin(qp **Node, minKey *interface{}, minVal *interface{}) bool { func removeMin(qp **Node, minKey *interface{}, minVal *interface{}) bool {
q := *qp q := *qp
if q.c[0] == nil { if q.Children[0] == nil {
*minKey = q.Key *minKey = q.Key
*minVal = q.Value *minVal = q.Value
if q.c[1] != nil { if q.Children[1] != nil {
q.c[1].p = q.p q.Children[1].Parent = q.Parent
} }
*qp = q.c[1] *qp = q.Children[1]
return true return true
} }
fix := removeMin(&q.c[0], minKey, minVal) fix := removeMin(&q.Children[0], minKey, minVal)
if fix { if fix {
return removeFix(1, qp) return removeFix(1, qp)
} }
@ -255,7 +255,7 @@ func putFix(c int8, t **Node) bool {
return false return false
} }
if s.c[(c+1)/2].b == c { if s.Children[(c+1)/2].b == c {
s = singlerot(c, s) s = singlerot(c, s)
} else { } else {
s = doublerot(c, s) s = doublerot(c, s)
@ -277,14 +277,14 @@ func removeFix(c int8, t **Node) bool {
} }
a := (c + 1) / 2 a := (c + 1) / 2
if s.c[a].b == 0 { if s.Children[a].b == 0 {
s = rotate(c, s) s = rotate(c, s)
s.b = -c s.b = -c
*t = s *t = s
return false return false
} }
if s.c[a].b == c { if s.Children[a].b == c {
s = singlerot(c, s) s = singlerot(c, s)
} else { } else {
s = doublerot(c, s) s = doublerot(c, s)
@ -302,8 +302,8 @@ func singlerot(c int8, s *Node) *Node {
func doublerot(c int8, s *Node) *Node { func doublerot(c int8, s *Node) *Node {
a := (c + 1) / 2 a := (c + 1) / 2
r := s.c[a] r := s.Children[a]
s.c[a] = rotate(-c, s.c[a]) s.Children[a] = rotate(-c, s.Children[a])
p := rotate(c, s) p := rotate(c, s)
switch { switch {
@ -324,14 +324,14 @@ func doublerot(c int8, s *Node) *Node {
func rotate(c int8, s *Node) *Node { func rotate(c int8, s *Node) *Node {
a := (c + 1) / 2 a := (c + 1) / 2
r := s.c[a] r := s.Children[a]
s.c[a] = r.c[a^1] s.Children[a] = r.Children[a^1]
if s.c[a] != nil { if s.Children[a] != nil {
s.c[a].p = s s.Children[a].Parent = s
} }
r.c[a^1] = s r.Children[a^1] = s
r.p = s.p r.Parent = s.Parent
s.p = r s.Parent = r
return r return r
} }
@ -393,7 +393,7 @@ func (t *Tree) bottom(d int) *Node {
return nil return nil
} }
for c := n.c[d]; c != nil; c = n.c[d] { for c := n.Children[d]; c != nil; c = n.Children[d] {
n = c n = c
} }
return n return n
@ -416,18 +416,18 @@ func (n *Node) walk1(a int) *Node {
return nil return nil
} }
if n.c[a] != nil { if n.Children[a] != nil {
n = n.c[a] n = n.Children[a]
for n.c[a^1] != nil { for n.Children[a^1] != nil {
n = n.c[a^1] n = n.Children[a^1]
} }
return n return n
} }
p := n.p p := n.Parent
for p != nil && p.c[a] == n { for p != nil && p.Children[a] == n {
n = p n = p
p = p.p p = p.Parent
} }
return p return p
} }
@ -446,14 +446,14 @@ func (n *Node) String() string {
} }
func output(node *Node, prefix string, isTail bool, str *string) { func output(node *Node, prefix string, isTail bool, str *string) {
if node.c[0] != nil { if node.Children[0] != nil {
newPrefix := prefix newPrefix := prefix
if isTail { if isTail {
newPrefix += "│ " newPrefix += "│ "
} else { } else {
newPrefix += " " newPrefix += " "
} }
output(node.c[0], newPrefix, false, str) output(node.Children[0], newPrefix, false, str)
} }
*str += prefix *str += prefix
if isTail { if isTail {
@ -462,13 +462,13 @@ func output(node *Node, prefix string, isTail bool, str *string) {
*str += "┌── " *str += "┌── "
} }
*str += node.String() + "\n" *str += node.String() + "\n"
if node.c[1] != nil { if node.Children[1] != nil {
newPrefix := prefix newPrefix := prefix
if isTail { if isTail {
newPrefix += " " newPrefix += " "
} else { } else {
newPrefix += "│ " newPrefix += "│ "
} }
output(node.c[1], newPrefix, true, str) output(node.Children[1], newPrefix, true, str)
} }
} }

Loading…
Cancel
Save