Balanced augmented interval trees

November 10, 2019

Wrote a Go implementation of an interval tree. This particular one is using the augmented style, as outlined here. It is also balanced. The balancing logic is mostly taken from here, and adapted. I’ve highlighted the parts where the balancing logic is. You may remove those parts and still have a working, unbalanced, interval tree.

These are the structs we have, a AITree struct to represent the tree itself, and AITreeNode to represent the nodes in the tree.

type AITree struct {
	root *AITreeNode
}

type AITreeNode struct {
	low   int64
	high  int64
	max   int64
	value string
	left  *AITreeNode
	right *AITreeNode
	bal   int64}

The public interfaces of the tree. I’ve only implemented SetInterval, which adds an interval into the tree, and GetAllIntervals, which returns all values of intervals that contain the given point p.

func (t *AITree) GetAllIntervals(p int64) []string {
	if t.root != nil {
		return t.root.walk(p)
	} else {
		return []string{}
	}
}

func (t *AITree) SetInterval(low, high int64, value string) {
	node := &AITreeNode{
		low:   low,
		high:  high,
		max:   high,
		value: value,
	}
	if t.root == nil {
		node.max = high
		t.root = node
	} else {
		t.root.set(node)
		if t.root.bal < -1 || t.root.bal > 1 {			t.rebalance()		}	}
}

walk traverses the tree, adding eligible intervals to the result.

func (n *AITreeNode) walk(p int64) []string {
	var output []string
	if p > n.max {
		return output
	}
	if n.left != nil {
		output = append(output, n.left.walk(p)...)
	}
	if (n.low <= p) && (p <= n.high) {
		output = append(output, n.value)
	}
	if p < n.low {
		return output
	}
	if n.right != nil {
		output = append(output, n.right.walk(p)...)
	}
	return output
}

From here on is mostly the balancing logic. We use a temporary sentinel parent for the purposes of rebalancing the root node itself.

The set method returns a boolean indicating true if the height of the tree rooted at n has increased as a result of setting the new node t.

func (t *AITree) rebalance() {	sentinel := &AITreeNode{left: t.root}	sentinel.rebalance(t.root)	t.root = sentinel.left}
func (n *AITreeNode) set(t *AITreeNode) bool {
	if t.high > n.max {
		n.max = t.high
	}
	if t.low <= n.low {
		if n.left == nil {
			n.left = t
			if n.right == nil {				n.bal = -1			} else {				n.bal = 0			}		} else {
			if n.left.set(t) {
				if n.left.bal < -1 || n.left.bal > 1 {					n.rebalance(n.left)				} else {					n.bal--				}			}		}
	} else {
		if n.right == nil {
			n.right = t
			if n.left == nil {				n.bal = -1			} else {				n.bal = 0			}		} else {
			if n.right.set(t) {
				if n.right.bal < -1 || n.right.bal > 1 {					n.rebalance(n.right)				} else {					n.bal++				}			}		}
	}
	return n.bal != 0
}

func (n *AITreeNode) rebalance(t *AITreeNode) {	if t.bal == -2 && t.left.bal == -1 {		n.rotateRight(t)	} else if t.bal == 2 && t.right.bal == 1 {		n.rotateLeft(t)	} else if t.bal == -2 && t.left.bal == 1 {		n.rotateLeftRight(t)	} else if t.bal == 2 && t.right.bal == -1 {		n.rotateRightLeft(t)	}}
func (n *AITreeNode) rotateLeft(c *AITreeNode) {	r := c.right	c.right = r.left	r.left = c	if c == n.left {		n.left = r	} else {		n.right = r	}	c.bal = 0	r.bal = 0}
func (n *AITreeNode) rotateRight(c *AITreeNode) {	l := c.left	c.left = l.right	l.right = c	if c == n.left {		n.left = l	} else {		n.right = l	}	c.bal = 0	l.bal = 0}
func (n *AITreeNode) rotateRightLeft(c *AITreeNode) {	c.right.left.bal = 1	c.rotateRight(c.right)	c.right.bal = 1	n.rotateLeft(c)}
func (n *AITreeNode) rotateLeftRight(c *AITreeNode) {	c.left.right.bal = -1	c.rotateLeft(c.left)	c.left.bal = -1	n.rotateRight(c)}