Keep It Balanced

const std = @import("std");

const Tree = struct {
    root: ?*Node = null,
    len: usize = 0,
    allocator: std.mem.Allocator,
    /// The program counts rotations. The output shows that the balance is
    /// cheap: a few rotations for a full tree, not one per node.
    rotations: usize = 0,

    const Node = struct {
        value: i32,
        left: ?*Node = null,
        right: ?*Node = null,
        /// A leaf has height 1. The node stores the height, and the code does
        /// not compute it again. To compute it again is an O(n) walk, and this
        /// structure exists to prevent that walk.
        height: i32 = 1,
    };

    fn init(allocator: std.mem.Allocator) Tree {
        return .{ .allocator = allocator };
    }

    fn deinit(self: *Tree) void {
        freeSubtree(self.allocator, self.root);
        self.root = null;
        self.len = 0;
    }

    fn freeSubtree(allocator: std.mem.Allocator, maybe: ?*Node) void {
        const node = maybe orelse return;
        freeSubtree(allocator, node.left);
        freeSubtree(allocator, node.right);
        allocator.destroy(node);
    }

    fn heightOf(maybe: ?*Node) i32 {
        return if (maybe) |node| node.height else 0;
    }

    fn refresh(node: *Node) void {
        node.height = 1 + @max(heightOf(node.left), heightOf(node.right));
    }

    /// Left height minus right height. The invariant is that this value stays
    /// in -1..1 for each node in the tree. Two means that one side is a full
    /// level deeper than the other side, and the node must rotate.
    fn balance(maybe: ?*Node) i32 {
        const node = maybe orelse return 0;
        return heightOf(node.left) - heightOf(node.right);
    }

    /// Rotate right: the left child becomes the new root of this subtree.
    ///
    ///       y            x
    ///      / \          / \
    ///     x   C   ->   A   y
    ///    / \              / \
    ///   A   B            B   C
    ///
    /// B gets a different parent, and nothing else moves. The in-order
    /// sequence is A x B y C before and after. Thus a rotation keeps the
    /// search property with no extra work.
    fn rotateRight(self: *Tree, y: *Node) *Node {
        const x = y.left.?;
        y.left = x.right;
        x.right = y;
        refresh(y);
        refresh(x);
        self.rotations += 1;
        return x;
    }

    fn rotateLeft(self: *Tree, x: *Node) *Node {
        const y = x.right.?;
        x.right = y.left;
        y.left = x;
        refresh(x);
        refresh(y);
        self.rotations += 1;
        return y;
    }

    fn insert(self: *Tree, value: i32) !void {
        self.root = try self.insertInto(self.root, value);
    }

    /// Recursive, because the rebalance occurs on the path back *up*. On the
    /// path from the insertion point to the root, the code updates the height
    /// of each node and checks its balance. The function returns the root of
    /// the subtree, which can be a new root. That return links a rotation into
    /// its parent.
    fn insertInto(self: *Tree, maybe: ?*Node, value: i32) !?*Node {
        const node = maybe orelse {
            const fresh = try self.allocator.create(Node);
            fresh.* = .{ .value = value };
            self.len += 1;
            return fresh;
        };

        if (value < node.value) {
            node.left = try self.insertInto(node.left, value);
        } else if (value > node.value) {
            node.right = try self.insertInto(node.right, value);
        } else {
            return node; // already present
        }

        refresh(node);

        // Four cases. The outer test tells which side is too deep. The inner
        // test tells where the new node went on that side. On the outside,
        // do one rotation. On the inside, rotate the child, then this node.
        const bias = balance(node);

        if (bias > 1 and value < node.left.?.value) {
            return self.rotateRight(node); // left-left
        }
        if (bias < -1 and value > node.right.?.value) {
            return self.rotateLeft(node); // right-right
        }
        if (bias > 1) { // left-right
            node.left = self.rotateLeft(node.left.?);
            return self.rotateRight(node);
        }
        if (bias < -1) { // right-left
            node.right = self.rotateRight(node.right.?);
            return self.rotateLeft(node);
        }
        return node;
    }

    fn contains(self: *const Tree, value: i32) bool {
        var current = self.root;
        while (current) |node| {
            if (value == node.value) return true;
            current = if (value < node.value) node.left else node.right;
        }
        return false;
    }

    fn writeInOrder(self: *const Tree, out: *std.Io.Writer) !void {
        try walk(out, self.root);
        try out.writeByte('\n');
    }

    fn walk(out: *std.Io.Writer, maybe: ?*Node) !void {
        const node = maybe orelse return;
        try walk(out, node.left);
        try out.print("{d} ", .{node.value});
        try walk(out, node.right);
    }

    /// Check the invariant at each node, and do not trust the insert code.
    /// The chapter is about this invariant. Thus the program checks it, and
    /// does not only claim it.
    fn isBalanced(maybe: ?*Node) bool {
        const node = maybe orelse return true;
        if (@abs(balance(node)) > 1) return false;
        return isBalanced(node.left) and isBalanced(node.right);
    }

    fn writeShape(self: *const Tree, out: *std.Io.Writer) !void {
        try shape(out, self.root, 0);
    }

    fn shape(out: *std.Io.Writer, maybe: ?*Node, depth: usize) !void {
        const node = maybe orelse return;
        try shape(out, node.right, depth + 1);
        for (0..depth) |_| try out.writeAll("    ");
        try out.print("{d}\n", .{node.value});
        try shape(out, node.left, depth + 1);
    }
};

pub fn main(init: std.process.Init) !void {
    var buf: [8192]u8 = undefined;
    var file_writer = std.Io.File.stdout().writerStreaming(init.io, &buf);
    const out = &file_writer.interface;

    var safe: std.heap.SafeAllocator = .init(std.heap.page_allocator, .{});
    defer std.debug.assert(safe.deinit() == 0);
    const allocator = safe.allocator();

    // The input that made the plain tree into one long chain: 1 to 15, in order.
    var tree = Tree.init(allocator);
    defer tree.deinit();

    try out.writeAll("inserting 1..15 in order, height after each:\n ");
    for (1..16) |i| {
        try tree.insert(@intCast(i));
        try out.print("{d}", .{Tree.heightOf(tree.root)});
        if (i < 15) try out.writeAll(" ");
    }
    try out.writeByte('\n');

    try out.print("\nfinal height {d}, ideal {d}, rotations {d}\n", .{
        Tree.heightOf(tree.root),
        std.math.log2_int(usize, 15) + 1,
        tree.rotations,
    });
    try out.print("balanced everywhere: {}\n", .{Tree.isBalanced(tree.root)});

    try out.writeAll("\nin-order: ");
    try tree.writeInOrder(out);

    try out.writeAll("\nshape (rotate the page 90 degrees clockwise):\n");
    try tree.writeShape(out);

    // Each rotation keeps the in-order sequence. Thus a rebalanced tree gives
    // exactly the same answers as the tree that it replaced.
    var missing: usize = 0;
    for (1..16) |i| {
        if (!tree.contains(@intCast(i))) missing += 1;
    }
    try out.print("\nvalues 1..15 still present: {d} of 15\n", .{15 - missing});

    // A different input: strictly descending. It causes the mirror-image cases.
    var descending = Tree.init(allocator);
    defer descending.deinit();
    var i: i32 = 15;
    while (i >= 1) : (i -= 1) try descending.insert(i);
    try out.print("\ninserting 15..1: height {d}, rotations {d}, balanced {}\n", .{
        Tree.heightOf(descending.root),
        descending.rotations,
        Tree.isBalanced(descending.root),
    });

    // The zig-zag input. It needs the cases with two rotations.
    var zigzag = Tree.init(allocator);
    defer zigzag.deinit();
    for ([_]i32{ 1, 15, 2, 14, 3, 13, 4, 12, 5, 11, 6, 10, 7, 9, 8 }) |v| try zigzag.insert(v);
    try out.print("zig-zag input:   height {d}, rotations {d}, balanced {}\n", .{
        Tree.heightOf(zigzag.root),
        zigzag.rotations,
        Tree.isBalanced(zigzag.root),
    });

    try out.flush();
}