treap

treap.go

package treap

type node struct {
	key   int
	prior int
	left  *node
	right *node
}

func zig(p *node) *node {
	right := p.right
	p.right = right.left
	right.left = p
	return right
}

func zag(p *node) *node {
	left := p.left
	p.left = left.right
	left.right = p
	return left
}

func insert(p *node, key int, prior int) *node {
	if p == nil {
		p = &node{
			key:   key,
			prior: prior,
		}
	} else if key < p.key {
		p.left = insert(p.left, key, prior)
		if p.left.prior < p.prior {
			p = zag(p)
		}
	} else {
		p.right = insert(p.right, key, prior)
		if p.right.prior < p.prior {
			p = zig(p)
		}
	}
	return p
}

func remove(p *node, key int) *node {
	if p == nil {
		return nil
	}
	if key < p.key {
		p.left = remove(p.left, key)
	} else if key > p.key {
		p.right = remove(p.right, key)
	} else if p.left == nil {
		p = p.right
	} else if p.right == nil {
		p = p.left
	} else if p.left.prior < p.right.prior {
		p = zag(p)
		p.right = remove(p.right, key)
	} else {
		p = zig(p)
		p.left = remove(p.left, key)
	}
	return p
}