An example AVL Tree (a self-balancing binary search tree) in Go
Here's a complete, runnable example using an AVL Tree (a self-balancing binary search tree). It orders custom objects by their UpdatedAt timestamp and automatically rebalances after each insertion.
package main
import (
"fmt"
"time"
)
// Record represents your custom object with a timestamp field.
type Record struct {
ID string
Name string
UpdatedAt time.Time
}
// LessThan defines the ordering for the tree based on UpdatedAt.
// Uses ID as a tie-breaker to guarantee strict ordering when timestamps match.
func (a Record) LessThan(b Record) bool {
if a.UpdatedAt.Equal(b.UpdatedAt) {
return a.ID < b.ID // Stable sort guarantee
}
return a.UpdatedAt.Before(b.UpdatedAt)
}
// Node is the tree node holding a record and balancing metadata.
type Node struct {
Key Record
Height int
Left *Node
Right *Node
}
// AVLTree maintains the root and provides insertion operations.
type AVLTree struct {
Root *Node
}
// height returns the height of a node (0 for nil).
func height(n *Node) int {
if n == nil {
return 0
}
return n.Height
}
// getMax returns the larger of two integers.
func getMax(a, b int) int {
if a > b {
return a
}
return b
}
// maxHeight computes height of a node including itself.
func maxHeight(n *Node) int {
if n == nil {
return 0
}
return getMax(height(n.Left), height(n.Right)) + 1
}
// rotateRight performs a right rotation around node y.
func rotateRight(y *Node) *Node {
x := y.Left
T2 := x.Right
// Perform rotation
x.Right = y
y.Left = T2
// Update heights (bottom-up)
y.Height = getMax(height(y.Left), height(y.Right)) + 1
x.Height = maxHeight(x)
return x // new root of this subtree
}
// rotateLeft performs a left rotation around node x.
func rotateLeft(x *Node) *Node {
y := x.Right
T2 := y.Left
// Perform rotation
y.Left = x
x.Right = T2
// Update heights (bottom-up)
x.Height = maxHeight(x)
y.Height = getMax(height(y.Left), height(y.Right)) + 1
return y // new root of this subtree
}
// Insert adds a record to the tree, maintaining AVL balance.
func (t *AVLTree) Insert(record Record) {
t.Root = insert(t.Root, record)
}
func insert(node *Node, record Record) *Node {
// 1. Standard BST insertion
if node == nil {
return &Node{Key: record, Height: 1}
}
if record.LessThan(node.Key) {
node.Left = insert(node.Left, record)
} else if !record.LessThan(node.Key) && !node.Key.LessThan(record) {
// Duplicate timestamp + ID: ignore (or update logic here if needed)
return node
} else {
node.Right = insert(node.Right, record)
}
// 2. Update height of current node
node.Height = maxHeight(node)
// 3. Get balance factor to check if node became unbalanced
balance := height(node.Left) - height(node.Right)
// 4. Rebalance if unbalanced (4 cases)
// Left-Left Case
if balance > 1 && record.LessThan(node.Left.Key) {
return rotateRight(node)
}
// Right-Right Case
if balance < -1 && !record.LessThan(node.Right.Key) {
return rotateLeft(node)
}
// Left-Right Case
if balance > 1 && !record.LessThan(node.Left.Key) {
node.Left = rotateLeft(node.Left)
return rotateRight(node)
}
// Right-Left Case
if balance < -1 && record.LessThan(node.Right.Key) {
node.Right = rotateRight(node.Right)
return rotateLeft(node)
}
// Return unchanged node pointer (standard BST insertion result)
return node
}
// InOrderTraversal retrieves records sorted by UpdatedAt.
func (t *AVLTree) InOrder() []Record {
var result []Record
inorder(t.Root, &result)
return result
}
func inorder(node *Node, result *[]Record) {
if node == nil {
return
}
inorder(node.Left, result)
*result = append(*result, node.Key)
inorder(node.Right, result)
}
func main() {
tree := &AVLTree{}
// Intentionally insert in chronological order to trigger imbalances
records := []Record{
{ID: "1", Name: "Alpha", UpdatedAt: time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)},
{ID: "2", Name: "Beta", UpdatedAt: time.Date(2023, 1, 2, 0, 0, 0, 0, time.UTC)},
{ID: "3", Name: "Gamma", UpdatedAt: time.Date(2023, 1, 3, 0, 0, 0, 0, time.UTC)},
{ID: "4", Name: "Delta", UpdatedAt: time.Date(2023, 1, 4, 0, 0, 0, 0, time.UTC)},
{ID: "5", Name: "Epsilon", UpdatedAt: time.Date(2023, 1, 5, 0, 0, 0, 0, time.UTC)},
}
fmt.Println("Inserting records in chronological order...")
for _, r := range records {
tree.Insert(r)
fmt.Printf(" โ Inserted: %s (%s)\n", r.ID, r.UpdatedAt.Format(time.RFC3339))
}
fmt.Println("\nโ
In-order traversal (sorted by UpdatedAt):")
for _, r := range tree.InOrder() {
fmt.Printf(" [%s] %s | Updated: %s\n", r.ID, r.Name, r.UpdatedAt.Format(time.RFC3339))
}
// Insert a record that disrupts chronological order to test rebalancing
tree.Insert(Record{ID: "0", Name: "Zero", UpdatedAt: time.Date(2023, 1, 15, 0, 0, 0, 0, time.UTC)})
fmt.Println("\n๐ After inserting a later date (triggers rebalancing)...")
for _, r := range tree.InOrder() {
fmt.Printf(" [%s] %s | Updated: %s\n", r.ID, r.Name, r.UpdatedAt.Format(time.RFC3339))
}
}
๐ Key Design Choices
- AVL Tree: Guarantees
O(log n)insertion/search/delete by enforcing that the height difference between left and right subtrees is at most 1. - Custom Comparator:
Record.LessThan()drives the BST ordering. The tie-breaker (ID) ensures deterministic behavior when two records share the exact sameUpdatedAttimestamp. - Bottom-Up Height Updates: Heights are recalculated after every insertion and rotation. This is critical for correct balance factor computation.
- Four Rotation Cases:
- LL โ Right Rotate
- RR โ Left Rotate
- LR โ Left then Right
- RL โ Right then Left
โ ๏ธ Production Considerations
| Aspect | Recommendation |
|---|---|
| Concurrency | This implementation is not thread-safe. Wrap with sync.RWMutex or use a concurrent tree library if accessed by multiple goroutines. |
| Duplicates | Currently ignores exact duplicates. For real-world data, you might want to update the existing node if a newer UpdatedAt arrives for the same ID. |
| Performance | For massive datasets or disk-backed storage, consider a B-Tree (e.g., github.com/gho-oss/gotreeset/btree). |
| Standard Library | Go's sort package handles slices/arrays efficiently. Use trees when you need dynamic insertions/deletions with guaranteed logarithmic bounds. |
Run the code as-is to see how the tree automatically rebalances even when inserted in strictly sorted order. Let me know if you need Delete, Search, or a concurrent-safe variant!