sadewa/goadsl — pull request #1 by sadewa (2014-12-10 13:01:04+00:00)
Title: AVLTRee.go

// AVL Tree

package main

import "fmt"

type Node struct {
Data int
Height int
Left *Node
Right *Node
}

// Return data: O(1)
func (n *Node) GetData() int {
return n.Data
}

// Set data: O(1)
func (n *Node) SetData(data int) {
n.Data = data
}

// Return height: O(1)
func (n *Node) GetHeight() int {
return n.Height
}

// Set height: O(1)
func (n *Node) SetHeight() {
var leftheight,rightheight int
if n.HasLeft() { leftheight = n.GetLeft().GetHeight() }
if n.HasRight() { rightheight = n.GetRight().GetHeight() }
if leftheight > rightheight {
n.Height = leftheight+1
} else {
n.Height = rightheight+1
}
}

func (n *Node) BalanceFactor() int {
var leftheight,rightheight int
if n.HasLeft() { leftheight = n.GetLeft().GetHeight() }
if n.HasRight() { rightheight = n.GetRight().GetHeight() }
return leftheight-rightheight
}

// Return left child: O(1)
func (n *Node) GetLeft() *Node {
return n.Left
}

// Set left child: O(1)
func (n *Node) SetLeft(left *Node) {
n.Left = left
}

// Test whether node has left child: O(1)
func (n *Node) HasLeft() bool {
if n.Left == nil { return false }
return true
}

// Return right child: O(1)
func (n *Node) GetRight() *Node {
return n.Right
}

// Set right child: O(1)
func (n *Node) SetRight(right *Node) {
n.Right = right
}

// Test whether node has right child: O(1)
func (n *Node) HasRight() bool {
if n.Right == nil { return false }
return true
}

type AVLTree struct {
Root *Node
}

// Create a new AVL Tree: O(1)
func NewAVLTree() *AVLTree {
return &AVLTree{}
}

func (a *AVLTree) Search(data int) bool {
if a.Root == nil { return false }
current := a.Root
for {
if data < current.GetData() {
if current.HasLeft() {
current = current.GetLeft()
} else {
break
}
} else if data > current.GetData(){
if current.HasRight() {
current = current.GetRight()
} else {
break
}
} else {
return true
}
}
return false
}

func (a *AVLTree) SearchMin(current *Node) int {
if current.HasLeft() { return a.SearchMin(current.GetLeft()) }
return current.GetData()
}

func (a *AVLTree) RotateRight(current *Node) *Node {
left := current.GetLeft()
current.SetLeft(left.GetRight())
left.SetRight(current)
current.SetHeight()
return left
}

func (a *AVLTree) RotateLeft(current *Node) *Node {
right := current.GetRight()
current.SetRight(right.GetLeft())
right.SetLeft(current)
current.SetHeight()
return right
}

func (a *AVLTree) Rebalance(current *Node) *Node {
b := current.BalanceFactor()
if b > 1 {
if current.GetLeft().BalanceFactor() < 0 {
current.SetLeft(a.RotateLeft(current.GetLeft()))
}
current = a.RotateRight(current)
} else if b < -1 {
if current.GetRight().BalanceFactor() > 0 {
current.SetRight(a.RotateRight(current.GetRight()))
}
current = a.RotateLeft(current)
}
current.SetHeight()
return current
}

func (a *AVLTree) Insert(data int) {
newnode := &Node{data,1,nil,nil}
if a.Root == nil { a.Root = newnode; return; }
a.Root = a.insert(a.Root,newnode)
}

func (a *AVLTree) insert(current, newnode *Node) *Node {
if newnode.GetData() < current.GetData() {
if current.HasLeft() {
current.SetLeft(a.insert(current.GetLeft(),newnode))
} else {
current.SetLeft(newnode)
}
} else if newnode.GetData() > current.GetData() {
if current.HasRight() {
current.SetRight(a.insert(current.GetRight(),newnode))
} else {
current.SetRight(newnode)
}
} else {
return nil
}
current.SetHeight()
current = a.Rebalance(current)
return current
}

func (a *AVLTree) Delete(data int) {
if a.Root == nil { return }
if a.Root.GetData() == data {
auxroot := &Node{0,0,a.Root,nil}
a.delete(a.Root,auxroot,data)
a.Root = auxroot.GetLeft()
} else {
a.delete(a.Root,nil,data)
}
}

func (a *AVLTree) delete(current, parent *Node, data int) *Node {
if data < current.GetData() {
if current.GetLeft() == nil { return nil }
current.SetLeft(a.delete(current.GetLeft(),current,data))
} else if data > current.GetData() {
if current.GetRight() == nil { return nil }
current.SetRight(a.delete(current.GetRight(),current,data))
} else {
if current.HasLeft() && current.HasRight() {
current.SetData(a.SearchMin(current.GetRight()))
current.SetRight(a.delete(current.GetRight(),current,current.GetData()))
} else if parent.GetLeft() == current {
if current.HasLeft() {
parent.SetLeft(current.GetLeft())
return current.GetLeft()
} else {
parent.SetLeft(current.GetRight())
return current.GetRight()
}
} else if parent.GetRight() == current {
if current.HasRight() {
parent.SetRight(current.GetRight())
return current.GetRight()
} else {
parent.SetRight(current.GetLeft())
return current.GetLeft()
}
}
}
current.SetHeight()
current = a.Rebalance(current)
return current
}

// Print BST with BFS: O(n)
func (a *AVLTree) BFSPrinting() {
if a.Root == nil { return }
q := make([]*Node,0,0)
q = append(q,a.Root)
for len(q) != 0 {
front := q[0]
q = q[1:]
fmt.Print(front.GetData()," ")
if front.HasLeft() { q = append(q,front.GetLeft()) }
if front.HasRight() { q = append(q,front.GetRight()) }
}
fmt.Println()
}

func main() {
a := NewAVLTree()
for i := 1; i <= 10; i++ {
a.Insert(i)
}
a.BFSPrinting()
a.Delete(9)
a.BFSPrinting()
a.Delete(10)
a.BFSPrinting()
}