Balanced Augmented Interval Trees

November 10, 2019

Wrote a Go imple­men­ta­tion of an inter­val tree. This par­tic­u­lar one is using the aug­ment­ed style, as out­lined here. It is also bal­anced. The bal­anc­ing logic is mostly taken from here, and adapt­ed. I’ve high­light­ed the parts where the bal­anc­ing logic is. You may remove those parts and still have a work­ing, unbal­anced, inter­val tree.

These are the structs we have, a AITree struct to rep­re­sent the tree itself, and AITreeNode to rep­re­sent 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 inter­faces of the tree. I’ve only imple­ment­ed SetInterval, which adds an inter­val into the tree, and GetAllIntervals, which returns all values of inter­vals that con­tain 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 tra­vers­es the tree, adding eli­gi­ble inter­vals 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 bal­anc­ing logic. We use a tem­po­rary sen­tinel parent for the pur­pos­es of rebal­anc­ing the root node itself.

The set method returns a boolean indi­cat­ing true if the height of the tree rooted at n has increased as a result of set­ting 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)
}