# 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)
}``````