Memorization and Dynamic Programming

Posted on | In

This post is about how to come up with solutions applying dynamic programming when all I can think about is brutal force.

Let’s start from a simple problem.

I. Climbing stairs

70. Climbing Stairs: count ways to reach the top of a stair case with n steps, supposing you can take 1 or 2 steps at a time.

Let f(n) be the number of ways to reach the top. If you take 1 step for the last climb, there will be f(n-1) ways; if you take 2 steps for the last climb, there will be f(n-2) ways. So the problem can be broken down sub-problems like this: f(n) = f(n-1) + f(n-2).

This is easy to implemented using recursion.

1
2
3
4
int takeStairs(int n) {
    if (n <= 2) return 1;
    return takeStairs(n - 1) + takeStairs(n - 2);
}

However, the time complexity of this implementation is O(2^n). It is exponential, and it will cause stack overflow eventually.

If you execute the recursive function on the paper, you will find that there are many redundanct calculation. Take n = 6 for example, we have to calculate f(4) twice and f(3) three times.

1
2
3
4
5
6
7
        f(6)
      /     \
    f(5)    f(4)
   /  \    /   \
f(4) f(3) f(3) f(2)
 /
f(3) ...

We can introduce memorization here to reduce redundancy.

1
2
3
4
5
int takeStairs(int n, int[] mark) {
    if (mark[n] > 0) return mark[n];
    mark[n] = takeStairs(n - 1) + takeStairs(n - 2);
    return mark[n];
}

A little improvement, but it still depends on function stack, which might lead to stack overflow.
So we should get rid of recursion, but how?
As shown in the diagram above, the function is executed from top to bottom and then traced back from bottom to top. Thus, the solution is to have f(n) = f(n-1) + f(n-2) implemented from bottom to top.

1
2
3
4
5
6
7
8
9
int takeStairs(int n) {
    int[] count = new int[n + 1];
    count[1] = 1;
    count[2] = 1;
    for (int i = 3; i <= n; i++) {
        count[i] = count[i - 1] + count[i - 2];
    }
    return count[n];
}

This bottom-top implementation applies the algorithm called Dynamic Programming. Here is the formulation for dp:

1
2
3
Let f(n) be the number of ways to reach the top.
f(n) = f(n-1) + f(n-2), if n>2
f(n) = 1, if n=1 or n=2

The space complexity of the above implementation is O(n), but this can be reduced to O(1) by maintaining the values of count[i-1] and count[i-2] only.

1
2
3
4
5
6
7
8
9
10
11
int takeStairs(int n) {
    int cnt1 = 1;
    int cnt2 = 1;
    int cnt;
    for (int i = 3; i <= n; i++) {
        cnt = cnt1 + cnt2;
        cnt2 = cnt1;
        cnt1 = cnt;
    }
    return cnt;
}

II. Decode ways

91. Decode Ways: determine the total number of ways to decode an encoded message which contains digits only; a message is encoded by converting letters to numbers using this mapping ‘A’ -> 1, …, ‘Z’ -> 26.

It is easy to apply the backtracking algorithm to this problem. For a digit in the message, it can be decoded to a letter; or it can combine the digit next to it and then be decoded to a new letter. Thus, we have two choices when we handle message[i]. We can try the first choice and then try the second one when it backtracks.

1
2
3
4
5
6
7
8
9
10
int countWays(int i, String message, int len) {
    if (i == len) return 1;
    int count = 0;
    char ch = message.charAt(i);
    if (ch > '0') 
        count += countWays(i + 1, message, len);
    if (i <= len - 2 && (ch == '1' || ch == '2' && message.charAt(i + 1) <= '6'))
        count += countWays(i + 2, message, len);
    return count;
}

It is a solution of brutal force with pruning. The time complexity is O(2^n), supposing the length of the message is n. How can we optimize this solution by applying dynamic programming?

First, introduce memorization: let count[i] be the number of decode ways for message.substring(i). It is because in the backtracking implementation countWays(i, message, len) is called for more than once for the same i, which can be avoided by memorization.
Second, have it implemented from bottom to top: calculate count[i] in the order that i descends from len-1 to 0.

Thus, we have the formulation for dp:

1
2
3
4
Let count[i] be the number of decode ways for message.substring(i).
count[i] = count[i+1] if message[i] > '0'
         + count[i+2] if message[i] == '1' or message[i] == '2' and message[i+1] <= '6'
count[i] = 1 if i == len

Note that this formulation looks a lot like the implementation in the countWays function. Actually, the thinking process is similar, so it shouldn’t hard to come up with a dp solution from brutal force.

Because count[i] depends on count[i+1] and count[i+2] only, we don’t need to maintain an array for storing all the intermediate results.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
int count = 0;
int count1 = 1;
int count2 = 1;
for (int i = len - 1; i >= 0; i--) {
    char ch = message.charAt(i);
    count = 0;
    if (ch > '0') {
        count += count1;
    }
    if (i <= len - 2 && (ch == '1' || ch == '2' && message.charAt(i + 1) <= '6') {
        count += count2;
    }
    count2 = count1;
    count1 = count;
}
return count;

III. Paint house

There are a row of n houses, each house can be painted with one of the k colors.
The cost of painting each house with a certain color is different.
You have to paint all the houses such that no two adjacent houses have the same color.
The cost of painting each house with a certain color is represented by a n*k cost costs.
For example, costs[0][0] is the cost of painting house 0 with color 0; costs[1][2]is the cost of painting house 1 with color 2, and so on…
Find the minimum cost to paint all houses.

Here is an implementation of brutal force that applies backtracking. It is quite straightforward, so it is easy to come up with.

1
2
3
4
5
6
7
8
9
10
11
12
13
// i represents the ith house, and j represents the jth color
void paint(int i, int cost, int[] mark, int[][] costs) {
    if (i == n) {
        if (cost < maxCost) maxCost = cost;
        return;
    }
    for (int j = 0; j < k; j++) {
        if (i > 0 && mark[i - 1] == i) continue;
        mark[i] = j;
        paint(i + 1, cost + costs[i][j], mark, costs);
        mark[i] = -1;
    }
}

How can we optimize this solution? Let’s first locate where the double counting is in the above implmentation.
Suppose there are 5 houses and 4 colors. If we paint house 0 with color 1, we are going to paint house 1 using colors except color 1, we will call the paint function like this:

1
2
3
paint(2, cost + costs[1][0], mark, costs); // mark[1] = 0
paint(2, cost + costs[1][2], mark, costs); // mark[1] = 2
paint(2, cost + costs[1][3], mark, costs); // mark[1] = 3

If we paint hourse 0 with color 3, we are going to paint house 1 using colors except color 3, we will call the paint function like this (and most of them are duplicates):

1
2
3
paint(2, cost + costs[1][0], mark, costs); // mark[1] = 0
paint(2, cost + costs[1][1], mark, costs); // mark[1] = 1
paint(2, cost + costs[1][2], mark, costs); // mark[1] = 2

Thus, we need to memorize the mediate results to avoid double counting. We can adopt the bottom-up idea - the minimum cost of house i using color j depends on the minimum cost of house i+1 using colors except color j, so we have the formulation as follows:

1
2
3
Let cost[i][j] be the minimum cost of for house i, i+1, ... n-1, on the premise that house i is painted with color j.
cost[i][j] = min{cost[i+1][k]} + costs[i][j], s.t. k != j
cost[i][j] = costs[i][j], if i == n-1

1
2
3
4
5
6
7
8
9
10
11
int[][] cost = new int[n + 1][k];
for (int i = n - 1; i >= 0; i--) {
    for (int j = 0; j < k; j++) {
        int min = 0;
        for (int t = 0; t < k; t++) {
            if (t != j && (min == 0 || min > costs[i + 1][t]))
                min = costs[i + 1][t];
        }
        cost[i][j] = min + costs[i][j];
    }
}

The time complexity for this implementation is O(n * k * k).

Follow-up: can we solve it in O(n * k)?
The answer is yes. It is unnecessary to search the k cost value for house i+1 to get the minimum cost. Since all we need is the minimum result from a previous loop for i+1, why not just keep the minimum result? The minimum result is not valid if house i+1 use color j to get the minimum result and we want to try the same color for house i. But we can keep the most two smallest costs instead of one, and use the second smallest cost as the minimum result if house i and house i+1 use the same color.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
int preMin1 = 0;
int preMin2 = 0;
int preColor1 = -1;
for (int i = n - 1; i >= 0; i--) {
    for (int j = 0; j < k; j++) {
        int cost = costs[i][j];
        if (j == preColor1) {
            cost += preMin2;
        } else {
            cost += preMin1;
        }

        int min1 = 0;
        int min2 = 0;
        int color1 = -1;
        if (min1 == 0 || min1 > cost) {
            min2 = min1;
            min1 = cost;
            color1 = j;
        } else if (min2 == 0 || min2 > cost) {
            min2 = cost;
        }

        preMin1 = min1;
        preMin2 = min2;
        preColor1 = color1;
    }
}

Summary

It’s ok that we can only come up with a brutal-force solution at first. Write it down, and look for double counting. Then we can introduce memorization for intermediate results to avoid repetition works. Last but not the least, memorization requires the intermediate results are calculated from bottom to top.
If you compare the implementations of both solutions, they are essentially the same, but dynamic programming introduces memorization and bottom-up implementation for optimization.