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