From e3980e5b80c39c604c21ed7cbe56a737fd2d258e Mon Sep 17 00:00:00 2001 From: Benjamin Scher Purcell Date: Sat, 25 Feb 2017 14:29:26 -0600 Subject: [PATCH] add avl implementation --- trees/avl/LICENSE | 13 + trees/avl/avl.go | 425 +++++++++++++++++++++++++ trees/avl/avl_test.go | 711 ++++++++++++++++++++++++++++++++++++++++++ trees/avl/iterator.go | 111 +++++++ 4 files changed, 1260 insertions(+) create mode 100644 trees/avl/LICENSE create mode 100644 trees/avl/avl.go create mode 100644 trees/avl/avl_test.go create mode 100644 trees/avl/iterator.go diff --git a/trees/avl/LICENSE b/trees/avl/LICENSE new file mode 100644 index 0000000..985d32d --- /dev/null +++ b/trees/avl/LICENSE @@ -0,0 +1,13 @@ +Copyright (c) 2017 Benjamin Scher Purcell + +Permission to use, copy, modify, and distribute this software for any +purpose with or without fee is hereby granted, provided that the above +copyright notice and this permission notice appear in all copies. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. diff --git a/trees/avl/avl.go b/trees/avl/avl.go new file mode 100644 index 0000000..09b5181 --- /dev/null +++ b/trees/avl/avl.go @@ -0,0 +1,425 @@ +// Copyright (c) 2017, Benjamin Scher Purcell. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package avl implements an AVL balanced binary tree. +// +// Structure is not thread safe. +// +package avl + +import ( + "io/ioutil" + "log" + + "github.com/emirpasic/gods/trees" + "github.com/emirpasic/gods/utils" +) + +func assertTreeImplementation() { + var _ trees.Tree = new(Tree) +} + +var dbgLog = log.New(ioutil.Discard, "avl: ", log.LstdFlags) + +// Tree holds elements of the AVL tree. +type Tree struct { + Root *Node + size int + Comparator utils.Comparator +} + +// A Node holds an Ordered element of the AVL tree in +// the Val field. +type Node struct { + Key interface{} + Value interface{} + c [2]*Node + p *Node + b int8 +} + +// NewWith instantiates a red-black tree with the custom comparator. +func NewWith(comparator utils.Comparator) *Tree { + return &Tree{Comparator: comparator} +} + +// NewWithIntComparator instantiates a red-black tree with the IntComparator, i.e. keys are of type int. +func NewWithIntComparator() *Tree { + return &Tree{Comparator: utils.IntComparator} +} + +// NewWithStringComparator instantiates a red-black tree with the StringComparator, i.e. keys are of type string. +func NewWithStringComparator() *Tree { + return &Tree{Comparator: utils.StringComparator} +} + +// Size returns the number of elements stored in the tree. +func (t *Tree) Size() int { + return t.size +} + +// Empty returns true if tree does not contain any nodes. +func (t *Tree) Empty() bool { + return t.size == 0 +} + +// Clear removes all nodes from the tree. +func (t *Tree) Clear() { + t.Root = nil + t.size = 0 +} + +// Get looks up val and returns the matching element if +// it is found. +// +// Val's Less implementation must be able to handle +// comparisons to elements stored in this tree. +func (t *Tree) Get(key interface{}) (value interface{}, found bool) { + n := t.Root + for n != nil { + cmp := t.Comparator(key, n.Key) + switch { + case cmp < 0: + n = n.c[0] + case cmp == 0: + return n.Value, true + case cmp > 0: + n = n.c[1] + } + } + return nil, false +} + +// Floor Finds floor node of the input key, return the floor node or nil if no ceiling is found. +// Second return parameter is true if floor was found, otherwise false. +// +// Floor node is defined as the largest node that is smaller than or equal to the given node. +// A floor node may not be found, either because the tree is empty, or because +// all nodes in the tree is larger than the given node. +// +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (t *Tree) Floor(key interface{}) (floor *Node, found bool) { + found = false + n := t.Root + for n != nil { + c := t.Comparator(key, n.Key) + switch { + case c == 0: + return n, true + case c < 0: + n = n.c[0] + case c > 0: + floor, found = n, true + n = n.c[1] + } + } + if found { + return + } + return nil, false +} + +// Ceiling finds ceiling node of the input key, return the ceiling node or nil if no ceiling is found. +// Second return parameter is true if ceiling was found, otherwise false. +// +// Ceiling node is defined as the smallest node that is larger than or equal to the given node. +// A ceiling node may not be found, either because the tree is empty, or because +// all nodes in the tree is smaller than the given node. +// +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (t *Tree) Ceiling(key interface{}) (floor *Node, found bool) { + found = false + n := t.Root + for n != nil { + c := t.Comparator(key, n.Key) + switch { + case c == 0: + return n, true + case c < 0: + floor, found = n, true + n = n.c[0] + case c > 0: + n = n.c[1] + } + } + if found { + return + } + return nil, false +} + +// Put inserts node into the tree. +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (t *Tree) Put(key interface{}, value interface{}) { + var put func(*Node, **Node) bool + put = func(p *Node, qp **Node) bool { + q := *qp + if q == nil { + t.size++ + *qp = &Node{Key: key, Value: value, p: p} + return true + } + + c := t.Comparator(key, q.Key) + if c == 0 { + q.Key = key + q.Value = value + return false + } + + if c < 0 { + c = -1 + } else { + c = 1 + } + a := (c + 1) / 2 + var fix bool + fix = put(q, &q.c[a]) + if fix { + return putFix(int8(c), qp) + } + return false + } + + put(nil, &t.Root) +} + +// Remove remove the node from the tree by key. +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (t *Tree) Remove(key interface{}) { + var remove func(**Node) bool + remove = func(qp **Node) bool { + q := *qp + if q == nil { + return false + } + + c := t.Comparator(key, q.Key) + if c == 0 { + t.size-- + if q.c[1] == nil { + if q.c[0] != nil { + q.c[0].p = q.p + } + *qp = q.c[0] + return true + } + fix := removemin(&q.c[1], &q.Key, &q.Value) + if fix { + return removeFix(-1, qp) + } + return false + } + + if c < 0 { + c = -1 + } else { + c = 1 + } + a := (c + 1) / 2 + fix := remove(&q.c[a]) + if fix { + return removeFix(int8(-c), qp) + } + return false + } + + remove(&t.Root) +} + +func removemin(qp **Node, minKey *interface{}, minVal *interface{}) bool { + q := *qp + if q.c[0] == nil { + *minKey = q.Key + *minVal = q.Value + if q.c[1] != nil { + q.c[1].p = q.p + } + *qp = q.c[1] + return true + } + fix := removemin(&q.c[0], minKey, minVal) + if fix { + return removeFix(1, qp) + } + return false +} + +func putFix(c int8, t **Node) bool { + s := *t + if s.b == 0 { + s.b = c + return true + } + + if s.b == -c { + s.b = 0 + return false + } + + if s.c[(c+1)/2].b == c { + s = singlerot(c, s) + } else { + s = doublerot(c, s) + } + *t = s + return false +} + +func removeFix(c int8, t **Node) bool { + s := *t + if s.b == 0 { + s.b = c + return false + } + + if s.b == -c { + s.b = 0 + return true + } + + a := (c + 1) / 2 + if s.c[a].b == 0 { + s = rotate(c, s) + s.b = -c + *t = s + return false + } + + if s.c[a].b == c { + s = singlerot(c, s) + } else { + s = doublerot(c, s) + } + *t = s + return true +} + +func singlerot(c int8, s *Node) *Node { + dbgLog.Printf("singlerot: enter %p:%v %d\n", s, s, c) + s.b = 0 + s = rotate(c, s) + s.b = 0 + dbgLog.Printf("singlerot: exit %p:%v\n", s, s) + return s +} + +func doublerot(c int8, s *Node) *Node { + dbgLog.Printf("doublerot: enter %p:%v %d\n", s, s, c) + a := (c + 1) / 2 + r := s.c[a] + s.c[a] = rotate(-c, s.c[a]) + p := rotate(c, s) + if r.p != p || s.p != p { + panic("doublerot: bad parents") + } + + switch { + default: + s.b = 0 + r.b = 0 + case p.b == c: + s.b = -c + r.b = 0 + case p.b == -c: + s.b = 0 + r.b = c + } + + p.b = 0 + dbgLog.Printf("doublerot: exit %p:%v\n", s, s) + return p +} + +func rotate(c int8, s *Node) *Node { + dbgLog.Printf("rotate: enter %p:%v %d\n", s, s, c) + a := (c + 1) / 2 + r := s.c[a] + s.c[a] = r.c[a^1] + if s.c[a] != nil { + s.c[a].p = s + } + r.c[a^1] = s + r.p = s.p + s.p = r + dbgLog.Printf("rotate: exit %p:%v\n", r, r) + return r +} + +// Keys returns all keys in-order +func (t *Tree) Keys() []interface{} { + keys := make([]interface{}, t.size) + it := t.Iterator() + for i := 0; it.Next(); i++ { + keys[i] = it.Key() + } + return keys +} + +// Values returns all values in-order based on the key. +func (t *Tree) Values() []interface{} { + values := make([]interface{}, t.size) + it := t.Iterator() + for i := 0; it.Next(); i++ { + values[i] = it.Value() + } + return values +} + +// Left returns the minimum element of the AVL tree +// or nil if the tree is empty. +func (t *Tree) Left() *Node { + return t.bottom(0) +} + +// Right returns the maximum element of the AVL tree +// or nil if the tree is empty. +func (t *Tree) Right() *Node { + return t.bottom(1) +} + +func (t *Tree) bottom(d int) *Node { + n := t.Root + if n == nil { + return nil + } + + for c := n.c[d]; c != nil; c = n.c[d] { + n = c + } + return n +} + +// Prev returns the previous element in an inorder +// walk of the AVL tree. +func (n *Node) Prev() *Node { + return n.walk1(0) +} + +// Next returns the next element in an inorder +// walk of the AVL tree. +func (n *Node) Next() *Node { + return n.walk1(1) +} + +func (n *Node) walk1(a int) *Node { + if n == nil { + return nil + } + + if n.c[a] != nil { + n = n.c[a] + for n.c[a^1] != nil { + n = n.c[a^1] + } + return n + } + + p := n.p + for p != nil && p.c[a] == n { + n = p + p = p.p + } + return p +} diff --git a/trees/avl/avl_test.go b/trees/avl/avl_test.go new file mode 100644 index 0000000..94db9e4 --- /dev/null +++ b/trees/avl/avl_test.go @@ -0,0 +1,711 @@ +// Copyright (c) 2017, Benjamin Scher Purcell. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package avl + +import ( + "fmt" + "testing" +) + +func TestAVLPut(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + tree.Put(1, "a") //overwrite + + if actualValue := tree.Size(); actualValue != 7 { + t.Errorf("Got %v expected %v", actualValue, 7) + } + if actualValue, expectedValue := fmt.Sprintf("%d%d%d%d%d%d%d", tree.Keys()...), "1234567"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue, expectedValue := fmt.Sprintf("%s%s%s%s%s%s%s", tree.Values()...), "abcdefg"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + + tests1 := [][]interface{}{ + {1, "a", true}, + {2, "b", true}, + {3, "c", true}, + {4, "d", true}, + {5, "e", true}, + {6, "f", true}, + {7, "g", true}, + {8, nil, false}, + } + + for _, test := range tests1 { + // retrievals + actualValue, actualFound := tree.Get(test[0]) + if actualValue != test[1] || actualFound != test[2] { + t.Errorf("Got %v expected %v", actualValue, test[1]) + } + } +} + +func TestAVLRemove(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + tree.Put(1, "a") //overwrite + + tree.Remove(5) + tree.Remove(6) + tree.Remove(7) + tree.Remove(8) + tree.Remove(5) + + if actualValue, expectedValue := fmt.Sprintf("%d%d%d%d", tree.Keys()...), "1234"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue, expectedValue := fmt.Sprintf("%s%s%s%s", tree.Values()...), "abcd"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue, expectedValue := fmt.Sprintf("%s%s%s%s", tree.Values()...), "abcd"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue := tree.Size(); actualValue != 4 { + t.Errorf("Got %v expected %v", actualValue, 7) + } + + tests2 := [][]interface{}{ + {1, "a", true}, + {2, "b", true}, + {3, "c", true}, + {4, "d", true}, + {5, nil, false}, + {6, nil, false}, + {7, nil, false}, + {8, nil, false}, + } + + for _, test := range tests2 { + actualValue, actualFound := tree.Get(test[0]) + if actualValue != test[1] || actualFound != test[2] { + t.Errorf("Got %v expected %v", actualValue, test[1]) + } + } + + tree.Remove(1) + tree.Remove(4) + tree.Remove(2) + tree.Remove(3) + tree.Remove(2) + tree.Remove(2) + + if actualValue, expectedValue := fmt.Sprintf("%s", tree.Keys()), "[]"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue, expectedValue := fmt.Sprintf("%s", tree.Values()), "[]"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if empty, size := tree.Empty(), tree.Size(); empty != true || size != -0 { + t.Errorf("Got %v expected %v", empty, true) + } + +} + +func TestAVLLeftAndRight(t *testing.T) { + tree := NewWithIntComparator() + + if actualValue := tree.Left(); actualValue != nil { + t.Errorf("Got %v expected %v", actualValue, nil) + } + if actualValue := tree.Right(); actualValue != nil { + t.Errorf("Got %v expected %v", actualValue, nil) + } + + tree.Put(1, "a") + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") // overwrite + tree.Put(2, "b") + + if actualValue, expectedValue := fmt.Sprintf("%d", tree.Left().Key), "1"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue, expectedValue := fmt.Sprintf("%s", tree.Left().Value), "x"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + + if actualValue, expectedValue := fmt.Sprintf("%d", tree.Right().Key), "7"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if actualValue, expectedValue := fmt.Sprintf("%s", tree.Right().Value), "g"; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLCeilingAndFloor(t *testing.T) { + tree := NewWithIntComparator() + + if node, found := tree.Floor(0); node != nil || found { + t.Errorf("Got %v expected %v", node, "") + } + if node, found := tree.Ceiling(0); node != nil || found { + t.Errorf("Got %v expected %v", node, "") + } + + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + + if node, found := tree.Floor(4); node.Key != 4 || !found { + t.Errorf("Got %v expected %v", node.Key, 4) + } + if node, found := tree.Floor(0); node != nil || found { + t.Errorf("Got %v expected %v", node, "") + } + + if node, found := tree.Ceiling(4); node.Key != 4 || !found { + t.Errorf("Got %v expected %v", node.Key, 4) + } + if node, found := tree.Ceiling(8); node != nil || found { + t.Errorf("Got %v expected %v", node, "") + } +} + +func TestAVLIteratorNextOnEmpty(t *testing.T) { + tree := NewWithIntComparator() + it := tree.Iterator() + for it.Next() { + t.Errorf("Shouldn't iterate on empty tree") + } +} + +func TestAVLIteratorPrevOnEmpty(t *testing.T) { + tree := NewWithIntComparator() + it := tree.Iterator() + for it.Prev() { + t.Errorf("Shouldn't iterate on empty tree") + } +} + +func TestAVLIterator1Next(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + tree.Put(1, "a") //overwrite + // │ ┌── 7 + // └── 6 + // │ ┌── 5 + // └── 4 + // │ ┌── 3 + // └── 2 + // └── 1 + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + key := it.Key() + switch key { + case count: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator1Prev(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(5, "e") + tree.Put(6, "f") + tree.Put(7, "g") + tree.Put(3, "c") + tree.Put(4, "d") + tree.Put(1, "x") + tree.Put(2, "b") + tree.Put(1, "a") //overwrite + // │ ┌── 7 + // └── 6 + // │ ┌── 5 + // └── 4 + // │ ┌── 3 + // └── 2 + // └── 1 + it := tree.Iterator() + for it.Next() { + } + countDown := tree.size + for it.Prev() { + key := it.Key() + switch key { + case countDown: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + countDown-- + } + if actualValue, expectedValue := countDown, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator2Next(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + key := it.Key() + switch key { + case count: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator2Prev(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + for it.Next() { + } + countDown := tree.size + for it.Prev() { + key := it.Key() + switch key { + case countDown: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + countDown-- + } + if actualValue, expectedValue := countDown, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator3Next(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(1, "a") + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + key := it.Key() + switch key { + case count: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator3Prev(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(1, "a") + it := tree.Iterator() + for it.Next() { + } + countDown := tree.size + for it.Prev() { + key := it.Key() + switch key { + case countDown: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := key, countDown; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + countDown-- + } + if actualValue, expectedValue := countDown, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator4Next(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(13, 5) + tree.Put(8, 3) + tree.Put(17, 7) + tree.Put(1, 1) + tree.Put(11, 4) + tree.Put(15, 6) + tree.Put(25, 9) + tree.Put(6, 2) + tree.Put(22, 8) + tree.Put(27, 10) + // │ ┌── 27 + // │ ┌── 25 + // │ │ └── 22 + // │ ┌── 17 + // │ │ └── 15 + // └── 13 + // │ ┌── 11 + // └── 8 + // │ ┌── 6 + // └── 1 + it := tree.Iterator() + count := 0 + for it.Next() { + count++ + value := it.Value() + switch value { + case count: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + } + if actualValue, expectedValue := count, tree.Size(); actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIterator4Prev(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(13, 5) + tree.Put(8, 3) + tree.Put(17, 7) + tree.Put(1, 1) + tree.Put(11, 4) + tree.Put(15, 6) + tree.Put(25, 9) + tree.Put(6, 2) + tree.Put(22, 8) + tree.Put(27, 10) + // │ ┌── 27 + // │ ┌── 25 + // │ │ └── 22 + // │ ┌── 17 + // │ │ └── 15 + // └── 13 + // │ ┌── 11 + // └── 8 + // │ ┌── 6 + // └── 1 + it := tree.Iterator() + count := tree.Size() + for it.Next() { + } + for it.Prev() { + value := it.Value() + switch value { + case count: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + default: + if actualValue, expectedValue := value, count; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + } + count-- + } + if actualValue, expectedValue := count, 0; actualValue != expectedValue { + t.Errorf("Size different. Got %v expected %v", actualValue, expectedValue) + } +} + +func TestAVLIteratorBegin(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.Begin() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + for it.Next() { + } + + it.Begin() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.Next() + if key, value := it.Key(), it.Value(); key != 1 || value != "a" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 1, "a") + } +} + +func TestAVLIteratorEnd(t *testing.T) { + tree := NewWithIntComparator() + it := tree.Iterator() + + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.End() + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it.End() + if it.node != nil { + t.Errorf("Got %v expected %v", it.node, nil) + } + + it.Prev() + if key, value := it.Key(), it.Value(); key != 3 || value != "c" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 3, "c") + } +} + +func TestAVLIteratorFirst(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + if actualValue, expectedValue := it.First(), true; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if key, value := it.Key(), it.Value(); key != 1 || value != "a" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 1, "a") + } +} + +func TestAVLIteratorLast(t *testing.T) { + tree := NewWithIntComparator() + tree.Put(3, "c") + tree.Put(1, "a") + tree.Put(2, "b") + it := tree.Iterator() + if actualValue, expectedValue := it.Last(), true; actualValue != expectedValue { + t.Errorf("Got %v expected %v", actualValue, expectedValue) + } + if key, value := it.Key(), it.Value(); key != 3 || value != "c" { + t.Errorf("Got %v,%v expected %v,%v", key, value, 3, "c") + } +} + +func benchmarkGet(b *testing.B, tree *Tree, size int) { + for i := 0; i < b.N; i++ { + for n := 0; n < size; n++ { + tree.Get(n) + } + } +} + +func benchmarkPut(b *testing.B, tree *Tree, size int) { + for i := 0; i < b.N; i++ { + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + } +} + +func benchmarkRemove(b *testing.B, tree *Tree, size int) { + for i := 0; i < b.N; i++ { + for n := 0; n < size; n++ { + tree.Remove(n) + } + } +} + +func BenchmarkAVLGet100(b *testing.B) { + b.StopTimer() + size := 100 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkGet(b, tree, size) +} + +func BenchmarkAVLGet1000(b *testing.B) { + b.StopTimer() + size := 1000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkGet(b, tree, size) +} + +func BenchmarkAVLGet10000(b *testing.B) { + b.StopTimer() + size := 10000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkGet(b, tree, size) +} + +func BenchmarkAVLGet100000(b *testing.B) { + b.StopTimer() + size := 100000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkGet(b, tree, size) +} + +func BenchmarkAVLPut100(b *testing.B) { + b.StopTimer() + size := 100 + tree := NewWithIntComparator() + b.StartTimer() + benchmarkPut(b, tree, size) +} + +func BenchmarkAVLPut1000(b *testing.B) { + b.StopTimer() + size := 1000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkPut(b, tree, size) +} + +func BenchmarkAVLPut10000(b *testing.B) { + b.StopTimer() + size := 10000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkPut(b, tree, size) +} + +func BenchmarkAVLPut100000(b *testing.B) { + b.StopTimer() + size := 100000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkPut(b, tree, size) +} + +func BenchmarkAVLRemove100(b *testing.B) { + b.StopTimer() + size := 100 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkRemove(b, tree, size) +} + +func BenchmarkAVLRemove1000(b *testing.B) { + b.StopTimer() + size := 1000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkRemove(b, tree, size) +} + +func BenchmarkAVLRemove10000(b *testing.B) { + b.StopTimer() + size := 10000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkRemove(b, tree, size) +} + +func BenchmarkAVLRemove100000(b *testing.B) { + b.StopTimer() + size := 100000 + tree := NewWithIntComparator() + for n := 0; n < size; n++ { + tree.Put(n, struct{}{}) + } + b.StartTimer() + benchmarkRemove(b, tree, size) +} diff --git a/trees/avl/iterator.go b/trees/avl/iterator.go new file mode 100644 index 0000000..b2df5c6 --- /dev/null +++ b/trees/avl/iterator.go @@ -0,0 +1,111 @@ +// Copyright (c) 2017, Benjamin Scher Purcell. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package avl + +import "github.com/emirpasic/gods/containers" + +func assertIteratorImplementation() { + var _ containers.ReverseIteratorWithKey = (*Iterator)(nil) +} + +// Iterator holding the iterator's state +type Iterator struct { + tree *Tree + node *Node + position position +} + +type position byte + +const ( + begin, between, end position = 0, 1, 2 +) + +// Iterator returns a stateful iterator whose elements are key/value pairs. +func (tree *Tree) Iterator() Iterator { + return Iterator{tree: tree, node: nil, position: begin} +} + +// Next moves the iterator to the next element and returns true if there was a next element in the container. +// If Next() returns true, then next element's key and value can be retrieved by Key() and Value(). +// If Next() was called for the first time, then it will point the iterator to the first element if it exists. +// Modifies the state of the iterator. +func (iter *Iterator) Next() bool { + switch iter.position { + case begin: + iter.position = between + iter.node = iter.tree.Left() + case between: + iter.node = iter.node.Next() + } + + if iter.node == nil { + iter.position = end + return false + } + return true +} + +// Prev moves the iterator to the next element and returns true if there was a previous element in the container. +// If Prev() returns true, then next element's key and value can be retrieved by Key() and Value(). +// If Prev() was called for the first time, then it will point the iterator to the first element if it exists. +// Modifies the state of the iterator. +func (iter *Iterator) Prev() bool { + switch iter.position { + case end: + iter.position = between + iter.node = iter.tree.Right() + case between: + iter.node = iter.node.Prev() + } + + if iter.node == nil { + iter.position = begin + return false + } + return true +} + +// Value returns the current element's value. +// Does not modify the state of the iterator. +func (iter *Iterator) Value() interface{} { + return iter.node.Value +} + +// Key returns the current element's key. +// Does not modify the state of the iterator. +func (iter *Iterator) Key() interface{} { + return iter.node.Key +} + +// Begin resets the iterator to its initial state (one-before-first) +// Call Next() to fetch the first element if any. +func (iter *Iterator) Begin() { + iter.node = nil + iter.position = begin +} + +// End moves the iterator past the last element (one-past-the-end). +// Call Prev() to fetch the last element if any. +func (iter *Iterator) End() { + iter.node = nil + iter.position = end +} + +// First moves the iterator to the first element and returns true if there was a first element in the container. +// If First() returns true, then first element's key and value can be retrieved by Key() and Value(). +// Modifies the state of the iterator +func (iter *Iterator) First() bool { + iter.Begin() + return iter.Next() +} + +// Last moves the iterator to the last element and returns true if there was a last element in the container. +// If Last() returns true, then last element's key and value can be retrieved by Key() and Value(). +// Modifies the state of the iterator. +func (iter *Iterator) Last() bool { + iter.End() + return iter.Prev() +}