Skip to content

# Balanced Augmented Interval Trees

Posted on: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)
}``````