RosettaCodeData/Task/Priority-queue/Stata/priority-queue.stata

103 lines
2.0 KiB
Plaintext

mata
struct Node {
real scalar time
transmorphic data
}
class Heap {
public:
struct Node rowvector nodes
real scalar len
real scalar size
real scalar minHeap
void new()
void push()
void siftup()
void siftdown()
struct Node scalar pop()
real scalar empty()
real scalar compare()
}
real scalar Heap::compare(a,b) {
struct Node scalar left, right
left = nodes[a]
right = nodes[b]
return(minHeap ? left.time<right.time : left.time>right.time)
}
void Heap::new() {
len = 0
size = 4
nodes = Node(1,size)
minHeap = 1 // defaults to min heap
}
real scalar Heap::empty() {
return(len==0)
}
void Heap::siftdown(real scalar index) {
parent = index
while (parent*2 <= len) {
child = parent*2
if (child+1 <= len ? compare(child+1,child) : 0) {
child++
}
if (compare(child,parent)) {
nodes[(child,parent)] = nodes[(parent,child)]
parent = child
} else break
}
}
void Heap::siftup(real scalar index) {
child = index
while(child>1) {
parent = floor(child/2)
if (compare(parent,child)) {
break
}
nodes[(child,parent)] = nodes[(parent,child)]
temp = child
child = parent
parent = temp
}
}
void Heap::push (real scalar time, transmorphic data) {
if (len + 1 >= size) {
nodes = (nodes, nodes)
size = size*2
}
len++
nodes[len].time = time
nodes[len].data = data
siftup(len)
}
struct Node scalar Heap::pop () {
if (len==0) {
_error(3000,"empty heap")
}
len--
struct Node scalar newnode
newnode.time = nodes[1].time
newnode.data = nodes[1].data
if (len>0) {
nodes[1] = nodes[len+1]
siftdown(1)
}
return(newnode)
}
void testHeap(real scalar minHeap) {
class Heap scalar h
struct Node scalar node
h = Heap()
h.minHeap = minHeap
h.push(3, "Clear drains")
h.push(4, "Feed cat")
h.push(5, "Make tea")
h.push(1, "Solve RC tasks")
h.push(2, "Tax return")
while (!h.empty()) {
node = h.pop()
printf("%f -> %s\n", node.time, node.data)
}
}
testHeap(1)
testHeap(0)
end