Optimal Binary Search Tree

The problem, simply

You get sorted keys, for example 10, 12, 20, and for each key how many times people search for it, for example 34, 8, 50.

You must build a binary search tree from these keys: smaller keys go to the left of a node and bigger keys to the right. Finding a key at depth d (the top is depth 1) takes d steps. The total cost is, for every key, its search count times its depth. Give back the smallest total cost you can get.

A story to remember it

Picture it

A nicely balanced tree is not the best here. With 12 on top the cost is 8 x 1 + 34 x 2 + 50 x 2 = 176, because the popular key 20 sits one step lower.

Examples

Way 1: the simple way

Write a helper cost(i, j): the best cost for the keys from position i to position j.

1. If there are no keys (i > j), the cost is 0.
2. Try every key r in the group as the top.
3. The left group i..r-1 and the right group r+1..j hang one step lower. Their best costs come from the helper, and every key in the whole group pays one more step, so add the sum of all their searches once.
4. Keep the smallest total.

Why add the sum once? Before, each key's depth was counted inside its own small tree. Hanging that tree under a new top makes every key one step deeper, which costs exactly its search count one more time. The top key itself also pays its count once, so the whole group's sum covers everyone.

class Solution {
    public int optimalCost(int[] keys, int[] freq) {
        return cost(freq, 0, keys.length - 1);
    }

    private int cost(int[] freq, int i, int j) {
        if (i > j) return 0;                      // no keys, no cost
        int groupSum = 0;
        for (int k = i; k <= j; k++) groupSum += freq[k];
        int best = Integer.MAX_VALUE;
        for (int r = i; r <= j; r++) {            // try each key on top
            int total = cost(freq, i, r - 1) + cost(freq, r + 1, j) + groupSum;
            best = Math.min(best, total);
        }
        return best;
    }
}

How long it takes: each call splits into two smaller calls for every possible top. The number of calls grows like O(3^n). Extra room: the waiting helpers, O(n).

Way 2: the clever way

Every helper call is a stretch of keys from i to j. There are only about n * n / 2 stretches. Solve them all once, shortest first, and keep the answers in a table cost[i][j].

1. Stretches of one key: cost[i][i] = freq[i].
2. For length 2, then 3, and so on: for each stretch, try every key r as the top, and read the two smaller stretches from the table.
3. cost[i][j] = min over r of (cost[i][r-1] + cost[r+1][j]) + sum of freq[i..j].
4. The answer sits in the top right corner, cost[0][n-1].

A running total (prefix sums) gives the sum of any stretch in one step.

class Solution {
    public int optimalCost(int[] keys, int[] freq) {
        int n = keys.length;
        if (n == 0) return 0;
        int[] pre = new int[n + 1];                  // running totals of freq
        for (int i = 0; i < n; i++) pre[i + 1] = pre[i] + freq[i];

        int[][] cost = new int[n][n];
        for (int i = 0; i < n; i++) cost[i][i] = freq[i];   // one key alone

        for (int len = 2; len <= n; len++) {
            for (int i = 0; i + len - 1 < n; i++) {
                int j = i + len - 1;
                int groupSum = pre[j + 1] - pre[i];
                cost[i][j] = Integer.MAX_VALUE;
                for (int r = i; r <= j; r++) {            // try each key on top
                    int left = (r > i) ? cost[i][r - 1] : 0;
                    int right = (r < j) ? cost[r + 1][j] : 0;
                    cost[i][j] = Math.min(cost[i][j], left + right + groupSum);
                }
            }
        }
        return cost[0][n - 1];
    }
}

How long it takes: about n * n / 2 stretches, and each one tries up to n tops, so O(n^3). Extra room: the table, O(n^2).

Watch out for

• Forgetting to add the group sum. Without it, depth never gets counted and every answer looks too small.
• Fill by stretch length, not by row. cost[i][j] needs shorter stretches like cost[i+1][j], which a plain top to bottom row loop has not filled yet.
• Empty sides. When r == i or r == j, one side has no keys and costs 0. Do not read cost[i][i-1].
• Some versions also count failed searches (keys not in the tree). Then each gap between keys gets its own count too, but the table idea is the same.

What interviewers ask next

• Can you build the tree, not only the cost? Store the best top root[i][j] for every stretch, then build recursively: the top of 0..n-1, then the tops of its left and right stretches.
• Can you make it faster? Yes. Knuth noticed that the best top for i..j lies between the best tops of i..j-1 and i+1..j. Only trying those cuts the time to O(n^2).
• How is this like matrix chain multiplication? Both try every split point inside a stretch and combine two smaller stretches. They are the same family, interval DP.
• What if the searches change over time? Then a self adjusting tree, such as a splay tree, moves popular keys up as they get searched.

Remember this

Grown-up words: the signposts form a binary search tree, the search counts are frequencies, a stretch of keys is an interval, the table is interval dynamic programming, the running totals are prefix sums, and the faster trick is Knuth's optimization.