Convert a normal BST into a Balanced BST

The problem, simply

You get a binary search tree (BST) that leans badly to one side. It might even be a long stick. Give back a BST with the same numbers that is balanced: at every spot, the heights of the left side and the right side differ by at most 1.

Why it matters: searching a BST costs its height. A stick of a million numbers is a million steps tall. A balanced tree of a million numbers is only about 20 steps tall.

A story to remember it

Picture it

Examples

Way 1: the simple way

Lay the toys out in order, then hang the middle one first.

1. Read the tree left, me, right. In a BST that gives the spots in sorted order. Keep the spots themselves in a list, so no new spots are made.
2. build(lo, hi): if lo > hi, return null. Otherwise take the middle spot mid, give it build(lo, mid - 1) as its left side and build(mid + 1, hi) as its right side, and return it.

import java.util.*;

class Solution {
    public TreeNode balanceBST(TreeNode root) {
        List<TreeNode> spots = new ArrayList<>();
        inorder(root, spots);                      // the spots in sorted order
        return build(spots, 0, spots.size() - 1);
    }

    private void inorder(TreeNode node, List<TreeNode> spots) {
        if (node == null) return;
        inorder(node.left, spots);
        spots.add(node);
        inorder(node.right, spots);
    }

    private TreeNode build(List<TreeNode> spots, int lo, int hi) {
        if (lo > hi) return null;
        int mid = (lo + hi) / 2;                   // the middle toy hangs on top
        TreeNode node = spots.get(mid);
        node.left = build(spots, lo, mid - 1);
        node.right = build(spots, mid + 1, hi);
        return node;
    }
}

How long it takes: one read and one build, each touching every spot once, O(n). Extra room: the list of n spots, O(n). This is the answer most interviewers expect.

Way 2: the clever way

If you may not use a list, do it with only twists, called rotations. A rotation lifts a child above its parent and keeps the BST order. This is the Day, Stout and Warren method.

1. Make a stick. Walk down the right edge from a dummy spot above the top. Whenever a spot has a left child, rotate that child up. When there is no left child, step right. At the end every spot hangs to the right, in sorted order. Count them: n.
2. Fold the stick. Let full be the biggest number of the form 2^k - 1 that is at most n. First rotate left n - full times, at every second spot, so the extra spots drop down to be the bottom row.
3. Then for m = full / 2, full / 4, ... down to 1: walk down the right edge and rotate every second spot left, m times. Each round folds the stick in half and adds one row.

class Solution {
    public TreeNode balanceBST(TreeNode root) {
        TreeNode dummy = new TreeNode(0);          // a hook above the top
        dummy.right = root;
        int n = makeStick(dummy);
        int full = Integer.highestOneBit(n + 1) - 1; // biggest 2^k - 1 that is <= n
        fold(dummy, n - full);                     // drop the extra spots to the bottom row
        for (int m = full / 2; m > 0; m /= 2) fold(dummy, m);
        return dummy.right;
    }

    // rotate right until nothing has a left child; returns the number of spots
    private int makeStick(TreeNode dummy) {
        int n = 0;
        TreeNode tail = dummy, rest = dummy.right;
        while (rest != null) {
            if (rest.left != null) {
                TreeNode up = rest.left;           // the left child goes up
                rest.left = up.right;
                up.right = rest;
                tail.right = up;
                rest = up;
            } else {
                n++;
                tail = rest;
                rest = rest.right;
            }
        }
        return n;
    }

    // walk the right edge and rotate every second spot left, m times
    private void fold(TreeNode dummy, int m) {
        TreeNode scan = dummy;
        for (int i = 0; i < m; i++) {
            TreeNode down = scan.right;            // this spot goes down
            scan.right = down.right;               // its right child comes up
            scan = scan.right;
            down.right = scan.left;
            scan.left = down;
        }
    }
}

How long it takes: making the stick uses at most n rotations, and the folds use n/2 + n/4 + ..., less than n. So O(n). Extra room: a few fingers and a dummy, O(1).

Watch out for

• Forgetting to clear old links. In Way 1, build overwrites both left and right of every spot, so old links never survive. If you skip the null returns, leaves keep their old children.
• Using (lo + hi) / 2 with huge indices can overflow. lo + (hi - lo) / 2 is safe.
• "Balanced" does not mean "perfect". With 6 spots one leaf is missing, and that is fine.

What interviewers ask next

• Build a balanced BST from a sorted array (LeetCode 108). It is the build half of Way 1.
• Build one from a sorted linked list without copying to an array (LeetCode 109). Count the list, then build the left half first while a finger walks the list in order. O(n) time, O(log n) room.
• How do real libraries stay balanced? They rotate a little after every insert and delete (AVL trees, red black trees), so they never need a full rebuild. Java's TreeMap is a red black tree.
• Scapegoat trees do the opposite: let the tree go out of shape, and when one part gets too lopsided, rebuild just that part with Way 1. That is still O(log n) per operation on average.

Remember this

Grown-up words: the even mobile is a height-balanced BST, laying the toys out in order is an inorder traversal, hanging the middle first is divide and conquer, a twist that lifts a child is a tree rotation, the stick is a vine, and the fold method is the Day Stout Warren (DSW) algorithm.