diff --git a/README.md b/README.md index 96593e9..50da48a 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Classic algorithms and data structures implemented in Go. Not for production use * [Priority Queue](https://github.com/arnauddri/algorithms/tree/master/data-structures/priority-queue) [(wiki)](http://en.wikipedia.org/wiki/Priority_queue) * [Queue](https://github.com/arnauddri/algorithms/tree/master/data-structures/queue) [(wiki)](http://en.wikipedia.org/wiki/Queue_%28abstract_data_type%29) * [Stack](https://github.com/arnauddri/algorithms/tree/master/data-structures/stack) [(wiki)](http://en.wikipedia.org/wiki/Stack_%28abstract_data_type%29) +* [Treap](https://github.com/DarioBalinzo/algorithms/tree/master/data-structures/treap) [(wiki)](https://en.wikipedia.org/wiki/Treap) + #### Graph algorithms diff --git a/data-structures/treap/treap.go b/data-structures/treap/treap.go new file mode 100644 index 0000000..a1285b0 --- /dev/null +++ b/data-structures/treap/treap.go @@ -0,0 +1,175 @@ +package treap + +import ( + "math" + "math/rand" +) + +type Node struct { + Value int + Priority int + Parent *Node + Left *Node + Right *Node +} + +//every node has a random generated priority +func NewNode(i int) *Node { + return &Node{Value: i, Priority: rand.Intn(math.MaxUint16)} +} + +func (n *Node) Compare(m *Node) int { + if n.Value < m.Value { + return -1 + } else if n.Value > m.Value { + return 1 + } else { + return 0 + } +} + +func (n *Node) ComparePriority(m *Node) int { + if n.Priority < m.Priority { + return -1 + } else if n.Priority > m.Priority { + return 1 + } else { + return 0 + } +} + +type Treap struct { + Head *Node + Size int +} + +func NewTreap(n *Node) *Treap { + if n == nil { + return &Treap{} + } + return &Treap{Head: n, Size: 1} +} + +func rotateToRight(x *Node, y *Node) { + x_old := *x + y_old := *y + + //a ok + //b + y.Left = x_old.Right + if x_old.Right != nil { + x.Right.Parent = y + } + + if y_old.Parent != nil { + if y_old.Parent.Left == y { + y.Parent.Left = x + } else { + y.Parent.Right = x + } + } + + x.Parent = y_old.Parent + + //c ok + //y x + x.Right = y + y.Parent = x + +} + +func rotateToLeft(x *Node, y *Node) { + x_old := *x + y_old := *y + + //a ok + //b + y.Right = x_old.Left + if x_old.Left != nil { + x.Left.Parent = y + } + //c ok + //y x + + if y_old.Parent != nil { + if y_old.Parent.Left == y { + y.Parent.Left = x + } else { + y.Parent.Right = x + } + } + + x.Parent = y_old.Parent + + x.Left = y + y.Parent = x +} + +func (t *Treap) Insert(i int) { + n := NewNode(i) + if t.Head == nil { + t.Head = n + t.Size++ + return + } + + h := t.Head + //walk down with node value as a binary tree + for { + if n.Compare(h) == -1 { + if h.Left == nil { + h.Left = n + n.Parent = h + break + } else { + h = h.Left + } + } else { + if h.Right == nil { + h.Right = n + n.Parent = h + break + } else { + h = h.Right + } + } + } + + //walk up using random priority as a min heap + for { + if n.Parent == nil || n.ComparePriority(n.Parent) == 1 { + break + } else { + if n == n.Parent.Left { + rotateToRight(n, n.Parent) + } else { + rotateToLeft(n, n.Parent) + } + } + if n.Parent == nil { + t.Head = n + break + } + } + + t.Size++ +} + +func (t *Treap) Search(i int) *Node { + h := t.Head + n := &Node{Value: i} + + for h != nil { + switch h.Compare(n) { + case -1: + h = h.Right + case 1: + h = h.Left + case 0: + return h + default: + panic("Node not found") + } + } + panic("Node not found") +} diff --git a/data-structures/treap/treap_test.go b/data-structures/treap/treap_test.go new file mode 100644 index 0000000..07aa316 --- /dev/null +++ b/data-structures/treap/treap_test.go @@ -0,0 +1,48 @@ +package treap + +import ( + "testing" +) + +func TestTreap(t *testing.T) { + + n := NewNode(1) + m := NewNode(2) + + // Test compare + if n.Compare(m) != -1 || m.Compare(n) != 1 || n.Compare(n) != 0 { + t.Error() + } + + Treap := NewTreap(n) + + Treap.Insert(4) + Treap.Insert(2) + Treap.Insert(5) + Treap.Insert(3) + Treap.Insert(6) + + if Treap.Size != 6 { + t.Error() + } + + five := Treap.Search(5) + + if five.Value != 5 { + t.Error() + } + + // Treap.Delete(5) + // + // if Treap.Size != 5 { + // t.Error() + // } + // + // four := *Treap.Search(4) + // if four.Right.Value != 6 || + // four.Left.Value != 2 || + // four.Parent.Value != 1 { + // fmt.Println(*Treap.Search(4)) + // t.Error() + // } +}