Python Arena
Step 6 of 6

Challenge: the 0/1 knapsack

A delivery drone can lift only so many kilos. Each parcel waiting at the depot has a weight and a fee for delivering it, and a parcel either flies whole or stays behind: you can't take half of one, or the same one twice. That "all or nothing" is the "0/1" in the name. Which parcels should go up to earn the most without overloading the drone? This is the 0/1 knapsack (a knapsack is a backpack), and it's the best-known DP problem in interviews.

Greedy fails again. Take parcels of 2, 3, 5 and 6 kilos, with fees of 3, 6, 8 and 9, and a drone that lifts 7 kilos. The biggest fee first takes the 6 kilos (fee 9) and then has room for nothing else: 9. The best fee per kilo first takes the 3 kilos (2 per kilo), then the 2 kilos: 9 again. But 2 + 5 kilos make exactly 7 and earn 3 + 8 = 11. Trying every choice of parcels finds that, but n parcels can be chosen in 2ⁿ ways.

The table

best[i][c] holds the most you can earn with the first i items and a capacity of c. Item i sits at index i - 1, with weight w and value v, and as in the backtracking module there are two choices for it:

  • Leave it: then everything comes from the earlier items, best[i - 1][c].
  • Take it, if it fits (w <= c): it earns v, and the earlier items fill the c - w that's left as well as they can, v + best[i - 1][c - w].

best[i][c] is the bigger of the two, and with no items everything is 0:

def knapsack_table(weights, values, capacity):
    n = len(weights)
    best = [[0] * (capacity + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        w, v = weights[i - 1], values[i - 1]
        for c in range(capacity + 1):
            best[i][c] = best[i - 1][c]                  # leave item i
            if w <= c and v + best[i - 1][c - w] > best[i][c]:
                best[i][c] = v + best[i - 1][c - w]      # take item i
    return best


for row in knapsack_table([2, 3, 5, 6], [3, 6, 8, 9], 7):
    print(row)
Output
[0, 0, 0, 0, 0, 0, 0, 0]
[0, 0, 3, 3, 3, 3, 3, 3]
[0, 0, 3, 6, 6, 9, 9, 9]
[0, 0, 3, 6, 6, 9, 9, 11]
[0, 0, 3, 6, 6, 9, 9, 11]

The answer is the bottom-right entry, 11. To find which items, walk back up from it: where best[i][c] differs from the entry just above it, item i was taken, so take its weight off c; either way, go up a row. Here row 3's 11 differs from the 9 above it (the 5 kilos fly, leaving a capacity of 2), and in that column row 1's 3 differs from the 0 above it (the 2 kilos fly). The table has (n + 1) × (capacity + 1) entries, each O(1): O(n × capacity) time and memory.

One row is enough

A row of the table is built from the row just above it and nothing else. So one list can stand in for the whole table, overwritten item after item, as long as the inner loop runs through the capacities from the top down:

best = [0] * (capacity + 1)
for w, v in zip(weights, values):
    for c in range(capacity, w - 1, -1):
        best[c] = max(best[c], v + best[c - w])

zip(weights, values) walks through the two lists side by side, giving the pairs (weights[0], values[0]), (weights[1], values[1]) and so on, and for w, v in unpacks each pair into two variables.

Why downward? c - w is smaller than c, so on the way down the loop reaches best[c - w] only after best[c]. So best[c - w] is still untouched by this item: it holds what the earlier items alone can earn, as the row above would. Going up is the classic mistake. Then best[c - w] may already include this item, so the item is taken again and again, like a coin in the coin change step:

def one_row(weights, values, capacity, downward):
    best = [0] * (capacity + 1)
    for w, v in zip(weights, values):
        if downward:
            capacities = range(capacity, w - 1, -1)
        else:
            capacities = range(w, capacity + 1)
        for c in capacities:
            best[c] = max(best[c], v + best[c - w])
    return best[capacity]


print(one_row([3], [8], 9, True), one_row([3], [8], 9, False))
Output
8 24

One item of weight 3 and a capacity of 9: going down packs it once, going up three times. The one row needs only O(capacity) memory, for the same O(n × capacity) time, but it can't tell you which items were taken: that needs the whole table.

Which sums can you make?

The same shape answers "can some of these numbers add up to exactly t?", and in Python a set of the sums you can make says it directly. Start with {0} (take nothing), and for each number, add it to every sum so far:

sums = {0}
for x in [3, 5, 9]:
    sums |= {s + x for s in sums}
print(sorted(sums))
Output
[0, 3, 5, 8, 9, 12, 14, 17]

a | b makes the union of two sets: a set of everything that's in a, in b or in both. sums |= new puts every item of new into sums, the way += adds to a number. Python builds the whole right side first, so the comprehension adds x to the sums as they were before this number, and no sum gets x twice: each number is used at most once, the set's way of going downward. Keep only the sums up to t (if s + x <= t), and the set never holds more than t + 1 sums: O(n × t) time and O(t) memory, instead of 2ⁿ subsets.

When the table is too big

O(n × capacity) grows with the capacity's size, not with its number of digits: that's called pseudo-polynomial time. For ten items and a capacity of 10 billion, simply trying the 2¹⁰ = 1,024 subsets is far faster, so DP pays when the capacity or the target is modest, as in most interview problems.

Your task

Challenge 1 of 3

count_subsets(numbers, target) should count the subsets of the positive numbers that add up to exactly target, each number used at most once. It keeps one row, ways[s], the number of ways to make each sum, as the coin change step did. But it counts too many: count_subsets([3, 1, 4, 2, 5], 7) returns 13. Fix it, so that it returns 3 (3 + 4, 2 + 5 and 1 + 4 + 2):

3

Challenge 2

Write knapsack(weights, values, capacity): the most total value of items whose weights add up to at most capacity, each item taken at most once. weights[i] and values[i] belong to item i, and all of them are positive whole numbers. The tests have 200 items and a capacity of 1,000, where trying every subset never ends, so fill a table, or one row going down.

For example, knapsack([2, 3, 5, 6], [3, 6, 8, 9], 7) returns 11:

11

Challenge 3

Two friends split a pile of parcels to carry home, and numbers holds the parcels' weights. Every parcel goes to one of the two, and the loads should be as even as possible. Write min_difference(numbers): the smallest possible difference between the two loads.

The lighter load is at most half the total, so collect every load it could be, up to half the total, with a set as in the lesson. The heaviest of those, s, is the most even split: the other friend carries total - s, and the difference is total - 2 * s. The tests have 200 weights up to 100, where trying every split never ends.

For example, min_difference([8, 3, 5, 9, 2]) returns 1 (8 + 5 is 13, against 3 + 9 + 2 = 14):

1