Interval DP: burst balloons in Python

Why this matters

Many problems ask for the best order to do a set of steps, where each step changes what the next step can do. Dynamic programming (DP, a method that solves a big problem by solving smaller versions of it once and saving the answers) works well for these problems when you can cut the work into pieces that do not affect each other.

This lesson covers interval DP on a problem called burst balloons. An interval DP solves a problem for every range of positions in a list, starting with short ranges and building to the full range. The previous lesson in this module, matrix chain multiplication, uses the same shape. Here the trick is to choose a different balloon to think about, so the pieces become independent.

The problem

A list is Python's ordered collection of values, written with square brackets, such as [3, 1, 5, 8]. Each item has an index, which is its position counted from 0. The first item is at index 0.

You are given a list of balloons. Each balloon has a number painted on it. When you burst balloon i (the balloon at index i), you get coins equal to the product of three numbers: the value of the balloon to its left, the value of balloon i, and the value of the balloon to its right. Those left and right balloons are the neighbours that are still in the list at that moment. After the burst, balloon i is removed and the list closes the gap.

Take the list [3, 1, 5, 8]. Here is one order that gives 167 coins:

1. Burst the balloon with value 1. Its neighbours are 3 and 5. You get 3 times 1 times 5, which is 15. The list is now [3, 5, 8].
2. Burst the balloon with value 5. Its neighbours are 3 and 8. You get 3 times 5 times 8, which is 120. The list is now [3, 8].
3. Burst the balloon with value 3. It has no left neighbour, so the missing neighbour counts as 1. You get 1 times 3 times 8, which is 24. The list is now [8].
4. Burst the balloon with value 8. Both neighbours are missing, so you get 1 times 8 times 1, which is 8.

The total is 15 + 120 + 24 + 8 = 167. The task is to find the order that gives the most coins. The list can be a few hundred balloons long, so trying every order is too slow. For 4 balloons there are 24 orders. For 10 balloons there are 3,628,800 orders. For 20 balloons the number is far too large to try.

Why bursting first does not work

The natural first idea is to pick a balloon, burst it, and then solve the rest. This fails because the coins from a burst depend on the neighbours at that moment, and those neighbours change as other balloons burst.

Take the list [2, 3]. There are two orders.

• Burst 2 first. Its neighbours are missing (count as 1) and 3. You get 1 times 2 times 3, which is 6. Then burst 3 with neighbours 1 and 1, which gives 1 times 3 times 1, which is 3. Total: 9.
• Burst 3 first. Its neighbours are 2 and missing (count as 1). You get 2 times 3 times 1, which is 6. Then burst 2, which gives 1 times 2 times 1, which is 2. Total: 8.

The best total is 9, and it came from bursting the 2 first. If you choose a balloon first, you have to know what its neighbours will be after everything else has burst, and you cannot know that yet. Bursting first leaves you with a messy question.

Think about the last balloon instead

Turn the question around. Instead of asking which balloon goes first, ask which balloon goes last.

Look at a range of balloons between two walls. A wall is a balloon that is not burst inside this range. The walls stay in place while the inside balloons are burst. Call the walls l (left) and r (right). Now pick one balloon k inside the range to be the last one burst.

When k is the last balloon burst, everything between l and k and everything between k and r is already gone. That means at the moment of the final burst, the neighbours of k are exactly l and r. The coins from bursting k are always the value at l, times the value at k, times the value at r.

Before k is burst, the balloons on the left of k sit between the walls l and k. The balloons on the right of k sit between k and r. Since k is still standing as a wall for both sides, the two sides cannot affect each other. They are two smaller problems of the same kind.

This gives the DP. Let dp[l][r] be the most coins you can get from bursting every balloon strictly between positions l and r, where l and r are never burst during that work. Then:

dp[l][r] = the maximum, over every k strictly between l and r, of
           dp[l][k] + dp[k][r] + value[l] * value[k] * value[r]

If there are no balloons between l and r, then dp[l][r] is 0. This happens when r is l + 1.

Sentinels at both ends

The formula needs a value at l and at r for every range, including the full list. The first balloon has no left neighbour and the last balloon has no right neighbour. A sentinel is a value added to the list only to make the edges work. Here the sentinel is a 1 at each end.

Multiplying by 1 does not change a product, so the 1s act as missing neighbours. They also act as walls that are never burst, because the answer is about the balloons between the sentinels.

For [3, 1, 5, 8] the list becomes [1, 3, 1, 5, 8, 1]. In Python this is [1] + nums + [1], where + joins two lists into a new list. Positions run from 0 to 5. The real balloons are at positions 1 to 4, and the answer is dp[0][5], the range between the two sentinels.

Without sentinels, the code needs special cases for the first and last balloon. Those special cases are where most bugs come from. The sentinel version has one rule for every balloon.

Worked example 1: the list 2 and 3

Start with [2, 3]. With sentinels the list is [1, 2, 3, 1], at positions 0 to 3.

Ranges with nothing inside, such as dp[0][1] and dp[1][2], are 0.

For the range dp[0][2], the only balloon inside is at position 1, which has value 2. So k is 1:

dp[0][2] = dp[0][1] + dp[1][2] + value[0] * value[1] * value[2]
         = 0 + 0 + 1 * 2 * 3
         = 6

For the range dp[1][3], the only balloon inside is at position 2, with value 3:

dp[1][3] = dp[1][2] + dp[2][3] + value[1] * value[2] * value[3]
         = 0 + 0 + 2 * 3 * 1
         = 6

Now the full range dp[0][3]. Two choices for k:

• k is 1 (value 2). The coins are dp[0][1] + dp[1][3] + value[0] * value[1] * value[3], which is 0 + 6 + 1 * 2 * 1 = 8.
• k is 2 (value 3). The coins are dp[0][2] + dp[2][3] + value[0] * value[2] * value[3], which is 6 + 0 + 1 * 3 * 1 = 9.

The maximum is 9. That matches the order where the 2 is burst first, which is the order the earlier section found by hand. The DP reached the same answer by choosing the last balloon, not the first.

Worked example 2: the list 3, 1, 5, 8

Now the four-balloon example from the start. The list with sentinels is [1, 3, 1, 5, 8, 1]. The table below was printed by a Python program that fills dp with the formula above. Row l and column r show dp[l][r].

arr = [1, 3, 1, 5, 8, 1]
row 0:    0    0    3   30  159  167
row 1:    0    0    0   15  135  159
row 2:    0    0    0    0   40   48
row 3:    0    0    0    0    0   40
row 4:    0    0    0    0    0    0
row 5:    0    0    0    0    0    0
best last balloon for [0,5] is index 4 (value 8)
best last balloon for [0,4] is index 1 (value 3)
best last balloon for [1,4] is index 3 (value 5)

Read the bottom-right cell of the first row. dp[0][5] is 167, the answer. The choice that gave it is k equal to 4, the balloon with value 8. So the last balloon burst is the 8. Its coins are value[0] * value[4] * value[5], which is 1 * 8 * 1 = 8. The left part, dp[0][4], is 159. That is the 15 + 120 + 24 from the first three bursts in the order shown at the start.

The right part dp[4][5] is 0 because nothing is between positions 4 and 5.

Now check a smaller cell. dp[0][2] is 3. Its range holds only the balloon with value 3, between walls with values 1 and 1. That gives 1 * 3 * 1 = 3, which is what the table shows.

The table is filled in order of range length. Each cell needs cells for shorter ranges. That is the reason for the loop order in the code below.

The Python code

Save this as burst_balloons.py. The function max_coins takes a list of numbers and returns the most coins.

Run it with:

python3 burst_balloons.py

The output is:

167
9
10
0

A few lines need explaining.

• len(nums) returns the number of items in the list.
• [[0] * (n + 2) for _ in range(n + 2)] is a list comprehension. It is a short way to write a loop that builds a list. This line builds a grid of zeros with one row and one column for each position, including the two sentinels. Zero is the right value for empty ranges. Each row is a separate list, which matters for a bug covered below.
• range(2, n + 2) gives the numbers from 2 up to, but not including, n + 2. Python ranges always leave out the end value. The loop over length starts at 2 because a range of length 1 has no balloons inside it. Its value stays 0.
• The loop over left stops early, so that right stays inside the list.
• The loop over k only uses positions strictly between left and right. range(left + 1, right) stops before right, so the walls are never burst.
• Python integers have no size limit, so the product of three large numbers is exact. Languages such as Java need a long type for this. In Python, print(2 ** 100) prints 1267650600228229401496703205376.
• The if __name__ == "__main__": line means the four print lines run only when this file is run directly. If another file imports max_coins, the prints do not run.

The loop runs over three positions for each cell, so the time grows about as n cubed. For 300 balloons there are about 4.5 million steps for the triples l < k < r, which is quick for a program of this size. On the machine used to check this lesson, the function took about 0.45 seconds for 300 random balloons. Your time will differ. The memory is one grid of n + 2 by n + 2 numbers.

Checking the code against brute force

A brute force check tries every order and keeps the best. It is too slow for big inputs but correct, which makes it a good test for the DP on small inputs. The program below compares both methods on 300 random lists of up to 7 balloons.

Save it as check_balloons.py in the same folder as burst_balloons.py.

import random
 
from burst_balloons import max_coins
 
 
def brute(values):
    if not values:
        return 0
    best = None
    for i, value in enumerate(values):
        left = values[i - 1] if i > 0 else 1
        right = values[i + 1] if i + 1 < len(values) else 1
        rest = values[:i] + values[i + 1:]
        total = left * value * right + brute(rest)
        if best is None or total > best:
            best = total
    return best
 
 
def main():
    rnd = random.Random(42)
    checked = 0
    for _ in range(300):
        n = 1 + rnd.randrange(7)
        nums = [rnd.randrange(10) for _ in range(n)]
        if max_coins(nums) != brute(nums):
            print("MISMATCH", nums)
            return
        checked += 1
    print("random checks passed:", checked)
 
 
if __name__ == "__main__":
    main()

Run it with python3 check_balloons.py. The output is:

random checks passed: 300

Some parts of this program need explaining.

• random.Random(42) makes a random number generator that starts from a fixed seed. A seed is the starting number for the random sequence. With the same seed, the same random lists appear on every run.
• rnd.randrange(10) gives a whole number from 0 to 9.
• enumerate(values) gives each item together with its index, so for i, value in enumerate(values) loops over both.
• values[:i] + values[i + 1:] is a slice. A slice takes part of a list. This one copies the list without the item at index i.
• values[i - 1] alone would be wrong when i is 0. In Python, index -1 means the last item, so the if guard gives the missing neighbour the value 1 instead.

Common mistakes

Looping in the wrong order

The DP cell dp[l][r] needs dp[l][k] and dp[k][r] to be finished first. Both of those ranges are shorter than [l, r], so a loop by length handles them. A loop that goes over left from small to large and computes every right for each left can read cells that are not yet filled.

Here is the same formula with the loops in the wrong order:

def max_coins_wrong_order(nums):
    n = len(nums)
    arr = [1] + nums + [1]
    dp = [[0] * (n + 2) for _ in range(n + 2)]
    for left in range(n + 2):
        for right in range(left + 2, n + 2):
            best = 0
            for k in range(left + 1, right):
                coins = dp[left][k] + dp[k][right] + arr[left] * arr[k] * arr[right]
                best = max(best, coins)
            dp[left][right] = best
    return dp[0][n + 1]
 
 
print(max_coins_wrong_order([3, 1, 5, 8]))

Its output is:

63

The correct answer is 167. The buggy version reads dp[k][right] for k greater than left before those cells are filled, so it reads 0 and returns a wrong number. Nothing crashes and nothing tells you the answer is wrong, so this kind of bug is easy to miss. The brute force check above catches it.

Forgetting the sentinels

If you use the balloons directly and do not add the 1s, the code reads one position too far at the edge. Here is a small function that reads the right neighbour of a balloon, without any sentinel:

def right_neighbour(nums, i):
    return nums[i + 1]
 
 
print(right_neighbour([3, 1, 5, 8], 3))

Python stops the program with this error:

Traceback (most recent call last):
  File "/path/to/right_neighbour.py", line 5, in <module>
    print(right_neighbour([3, 1, 5, 8], 3))
          ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/path/to/right_neighbour.py", line 2, in right_neighbour
    return nums[i + 1]
           ~~~~^^^^^^^
IndexError: list index out of range

The path depends on where you saved the file. The last line says the program used an index that is not in the list. The list has indexes 0 to 3, so index 4 does not exist. The fix is to build arr with one extra item at each end, as the code above does.

Lists that share rows

A common Python mistake is to build the grid with * instead of a comprehension:

grid = [[0] * 3] * 3
grid[0][0] = 5
print(grid)

Its output is:

[[5, 0, 0], [5, 0, 0], [5, 0, 0]]

[[0] * 3] makes one row. Multiplying the outer list by 3 puts the same row object into all three places. Changing grid[0][0] changes the one shared row, so all three rows show the 5. The fix is the comprehension used in max_coins:

grid = [[0] * 3 for _ in range(3)]
grid[0][0] = 5
print(grid)

Its output is:

[[5, 0, 0], [0, 0, 0], [0, 0, 0]]

The comprehension runs [0] * 3 once for each row, so each row is a separate list.

Letting a wall be burst

The loop over k must start at left + 1 and stop before right. If you write range(left + 1, right + 1), the code can choose the wall itself as the last balloon. The formula then counts a burst whose neighbours are not the ones it assumes. Keep the bounds strict: left + 1 at the start, and the stop value right, which Python leaves out.

Removing by value instead of by position

The brute force check removes a balloon by slicing, which uses the position. The list method remove works differently. It removes the first item equal to the value you give, which may be a different balloon when values repeat:

values = [1, 2, 1]
values.remove(1)
print(values)

Its output is:

[2, 1]

The first 1 was removed, and the 1 at index 2 is still there. Use del values[i] when you mean a position.

Choosing the first balloon

A common wrong plan is to write dp as "burst balloon k first, then solve the left and right sides". The left and right sides then have balloons whose neighbours depend on k. The formula has to pick k as the last one, not the first.

Practice

Exercise 1

For the list [1, 5], find the best total by hand. Use the sentinels [1, 1, 5, 1], and check both orders.

Show the answer

Burst the 1 first. Its neighbours are the sentinel 1 and the 5, so the coins are 1 times 1 times 5, which is 5. Then burst the 5 with neighbours 1 and 1, which gives 1 times 5 times 1, which is 5. The total is 10.

Burst the 5 first. Its neighbours are 1 and the sentinel 1, which gives 5. Then burst the 1 with neighbours 1 and 1, which gives 1. The total is 6.

The best total is 10. The program also printed 10.

Exercise 2

Using the table from the worked example, what is dp[0][4] for [3, 1, 5, 8], and what does it mean?

Show the answer

It is 159. It is the best total for bursting the 3, the 1 and the 5 when the walls are the sentinel 1 at position 0 and the 8 at position 4. The 8 stays in place during that work, so it is a wall for the range.

Exercise 3

How many times in total does the inner loop run for a list with 100 balloons? The list has 102 positions including the sentinels.

Show the answer

The inner loop tries every position strictly between l and r. Across all ranges, the number of triples l < k < r among 102 positions is the number of ways to choose 3 positions from 102, which is 102 times 101 times 100 divided by 6. That equals 171,700 steps. Each step does a few multiplications and comparisons.

Exercise 4

Change the program so that it also returns the position of the last balloon burst for the full list. Describe the change in words and in code.

Show the answer

Keep a second grid, pick, built the same way as dp. Whenever a candidate coins is better than best, store k in pick[left][right]. After the loops, pick[0][n + 1] is the position of the last balloon in an optimal order. Its value is arr[pick[0][n + 1]]. The table from the worked example shows this value for [3, 1, 5, 8] is position 4, which holds the 8.

Exercise 5

Why must dp[l][r] treat l and r as walls that are not burst, even though those balloons may be burst in the full problem?

Show the answer

The formula looks only at the range between l and r. The values at l and r are the neighbours of every balloon in that range, so they must still be there. If a wall were burst inside the subproblem, its neighbours would change, and the formula's product would no longer be correct. The full problem is solved by the range that has the two sentinels as walls, and those sentinels are never burst.

Summary

• Thinking about the last balloon burst makes the left and right parts of a range independent.
• The formula is dp[l][r] = max over k of dp[l][k] + dp[k][r] + value[l] * value[k] * value[r].
• Add a 1 at each end of the list so every balloon has a left and right neighbour.
• Fill the table by range length, so smaller ranges are ready before the bigger ones that use them.
• Build each row of a grid with a comprehension, and keep the loop bounds strict so walls are never burst.
• Python integers do not overflow, so the answer can be any size. The time is about n cubed and the memory is about n squared. Check the code against brute force on small inputs.