[Coding Made Simple] Optimal Binary Search Tree

Given keys and frequency at which these keys are searched, how would you create a binary search tree from these keys 

such that the cost of searching is minimum.

The cost of searching is defined as the sum of all node's search frequency * its depth; Root node's depth is 1. 

Solution 1. Recursion

Algorithm:

1. first sort the input keys in ascending order based on their values;

2. from left to right of the sorted keys, pick keys[i] as the root node, then recursively calculate the search min cost of 

the left subtree and right subtree;

3. repeat step 2 for all keys and upate the depths of each node that gives the min search cost.

4. reconstruct the optimal BST using the minCostDepths[] that contains the depth of each node, this BST gives the min search cost.

This recursive solution is not efficient in that it has overlapping subprolems. This is demonstrated in the following example.

 Given keys {10, 11, 12, 13, 14}:

if 11 is used as root, then the subproblems of {10} and {12, 13, 14} are solved recursively;

if 12 is used as root, then the subproblems of {10, 11} and {13, 14} are solved recursively;

When solving subproblem of {12, 13, 14} in the case of 11 as root, we need to solve subproblem of {13, 14} after selecting 12 as root of subproblem {12, 13, 14};

This subproblem of {13, 14} is later solved again in the case of 12 as root; This is redundant work!

import java.util.Arrays;
import java.util.Comparator;

class TreeNode{
    int key;
    int freq;
    TreeNode left;
    TreeNode right;
    TreeNode(int key, int freq){
        this.key = key;
        this.freq = freq;
        this.left = null;
        this.right = null;
    }
}
public class OptimalBST {
    public TreeNode getOptimalBST(TreeNode[] nodes){
        if(nodes == null || nodes.length == 0){
            return null;
        }
        Comparator<TreeNode> comp = new Comparator<TreeNode>(){
            public int compare(TreeNode node1, TreeNode node2){
                return node1.key - node2.key;
            }
        };
        Arrays.sort(nodes, comp);
        int[] depths = new int[nodes.length];
        int[] minCostDepths = new int[nodes.length];
        getOptimalBSTRecursive(nodes, minCostDepths, depths, 0, nodes.length - 1, 1);
        return reconstructOptimalBST(nodes, minCostDepths, 0, nodes.length - 1);
    }
    private int getOptimalBSTRecursive(TreeNode[] nodes, int[] minCostDepths, int[] depths, int start, int end, int depth){
        if(start > end){
            return 0;
        }
        int minCost = Integer.MAX_VALUE;
        int cost = 0;
        for(int i = start; i <= end; i++){
            depths[i] = depth;
            cost += nodes[i].freq * depth;
            cost += getOptimalBSTRecursive(nodes, minCostDepths, depths, start, i - 1, depth + 1);
            cost += getOptimalBSTRecursive(nodes, minCostDepths, depths, i + 1, end, depth + 1);
            if(cost < minCost){
                minCost = cost;
                for(int j = start; j <= end; j++){
                    minCostDepths[j] = depths[j];
                }
            }
        }
        return minCost;
    }
    private TreeNode reconstructOptimalBST(TreeNode[] nodes, int[] minCostDepths, int start, int end){
        if(start > end){
            return null;
        }
        int minDepth = Integer.MAX_VALUE;
        int minDepthIdx = 0;
        for(int i = start; i <= end; i++){
            if(minCostDepths[i] < minDepth){
                minDepth = minCostDepths[i];
                minDepthIdx = i;
            }
        }
        nodes[minDepthIdx].left = reconstructOptimalBST(nodes, minCostDepths, start, minDepthIdx - 1);
        nodes[minDepthIdx].right = reconstructOptimalBST(nodes, minCostDepths, minDepthIdx + 1, end);
        return nodes[minDepthIdx];
    }
}

Solution 2. Dynamic Programming 

Given a list of n nodes that are sorted in ascending order using their keys,  nodes[0..... n - 1],

 State: T[i][j] contains the min cost of the bst using nodes[i....j] and the index of the root of this optimal bst.

Function:  For a given range of nodes, its min cost is the sum of each nodes' key 

T[i][j].cost = sum of nodes' freq in nodes[i....j] + min of {T[i][k - 1] + T[k + 1][j],  for k is from i to j}

Initialization: T[i][i].cost = nodes[i].freq; T[i][i].idx = i; 

Answer: T[0][n - 1] contains the min cost the optimal bst and its root idx.

The whole optimal bst can be constructed recursively base on the DP 2D array infos.

The DP solution avoids the duplicated work in solution 1 at the cost of O(n^2) extra memory. 

import java.util.Arrays;
import java.util.Comparator;

class BstNode {
    int key;
    int freq;
    BstNode left;
    BstNode right;
    BstNode(int key, int freq) {
        this.key = key;
        this.freq = freq;
        this.left = null;
        this.right = null;
    }
}
public class OptimalBST {
    class optimalBstEntry {
        int cost;
        int idx;
        optimalBstEntry(int c, int i) {
            this.cost = c;
            this.idx = i;
        }
    }
    public BstNode getOptimalBST(BstNode[] nodes) {
        if(nodes == null || nodes.length == 0) {
            return null;
        }
        Comparator<BstNode> comp = new Comparator<BstNode>() {
            public int compare(BstNode node1, BstNode node2) {
                return node1.key - node2.key;
            }
        };
        Arrays.sort(nodes, comp);
        optimalBstEntry[][] T = new optimalBstEntry[nodes.length][nodes.length];
        for(int i = 0; i < nodes.length; i++) {
            for(int j = i; j < nodes.length; j++)
            T[i][j] = new optimalBstEntry(Integer.MAX_VALUE, -1);
        }
        for(int i = 0; i < nodes.length; i++) {
            T[i][i].cost = nodes[i].freq;
            T[i][i].idx = i;
        }
        for(int len = 2; len <= nodes.length; len++) {
            for(int startIdx = 0; startIdx <= nodes.length - len; startIdx++) {
                int sum = 0; int min = Integer.MAX_VALUE;
                for(int i = startIdx; i < (startIdx + len); i++) {
                    sum += nodes[i].freq;
                }
                for(int rootIdx = startIdx; rootIdx <= (startIdx + len - 1); rootIdx++) {
                    int leftCost = 0, rightCost = 0;
                    if(rootIdx - 1 >= startIdx) {
                        leftCost = T[startIdx][rootIdx - 1].cost;
                    }
                    if(rootIdx + 1 <= startIdx + len - 1) {
                        rightCost = T[rootIdx + 1][startIdx + len - 1].cost;
                    }
                    if(leftCost + rightCost < min) {
                        T[startIdx][startIdx + len - 1].cost = sum + leftCost + rightCost;
                        T[startIdx][startIdx + len - 1].idx = rootIdx;
                        min = leftCost + rightCost;
                    }
                }
            }
        }
        //construct the optimal BST        
        return constructOptimalBst(nodes, T, 0, nodes.length - 1);
    }
    private BstNode constructOptimalBst(BstNode[] nodes, optimalBstEntry[][] T, int start, int end) {
        if(start > end) {
            return null;
        }
        int rootIdx = T[start][end].idx;
        BstNode root = nodes[rootIdx];
        root.left = constructOptimalBst(nodes, T, start, rootIdx - 1);
        root.right = constructOptimalBst(nodes, T, rootIdx + 1, end);
        return root;
    }
}

Further optimization

When calculating the sum of keys in a given range,  the above solution always start from stratch for a given range. 

For example, when calculating the sum of keys of the node 0, 1, 2, 3, it adds all 4 nodes' keys. Since we've already

calculated the key sums of 3 nodes, we can use prefix sum to memorize the previous calculation results. Doing this 

makes this sum operation from O(n) to O(1) since we only need to do one add operation for the newly included node.

However, this optimize does not change the asymptotic complexity, it only changes the constant factor. 

The runtime is still O(n^3).

Related Topics

Prefix Sum