Knapsack Variants: Groups, Two Budgets, Reconstruction and Counting

This is an extension topic beyond most interviews: it pays off in harder OAs where a knapsack comes with a twist, such as groups, a second budget, or a request for the actual chosen items. Everything here builds on the 0/1 recurrence f[j] = max(f[j], f[j - w] + v) with a descending loop.

Intuition

Each variant changes one thing about the basic model and adjusts the recurrence accordingly:

• Group knapsack: items come in groups and you may take at most one item per group.
• Two-dimensional cost: every item consumes two resources (for example time and money), so the state needs one dimension per budget.
• Reconstruction: report which items to take, not just the best value.
• Counting optimal solutions: how many subsets achieve the best value.
• Undoing an item: in a counting knapsack, remove one item's contribution in O(W) instead of recomputing.

When to use it in interviews / OAs

• "Pick at most one option from each category" or "take a prefix from each pile" means group knapsack.
• Two separate limits on the same selection (zeros and ones, weight and volume) mean a two-dimensional cost.
• "Return the chosen items" or "print the plan" requires a full table or choice records.
• "For every item, answer the question without that item" in a counting setting suggests undoing.

Group knapsack

The loop order guarantees at most one item per group: for each capacity j (descending), try every item of the group against values that still come from previous groups.

def group_knapsack(groups: list[list[tuple[int, int]]], cap: int) -> int:
    # groups[g] = list of (weight, value); weights are positive
    f = [0] * (cap + 1)  # f[j] = best value with weight <= j
    for group in groups:
        for j in range(cap, -1, -1):  # capacity loop OUTSIDE the item loop
            for w, v in group:
                if w <= j:
                    f[j] = max(f[j], f[j - w] + v)
    return f[cap]

If the item loop were outside, two items of the same group could both be added at different capacities, which turns it back into a plain 0/1 knapsack.

A typical application: from several stacks you may take coins only from the top, k coins in total. Taking the top t coins of a pile is "item t of that pile's group", with weight t and value equal to the prefix sum:

def max_value_of_coins(piles: list[list[int]], k: int) -> int:
    groups = []
    for pile in piles:
        prefix, items = 0, []
        for t, coin in enumerate(pile, 1):
            prefix += coin
            items.append((t, prefix))  # take the top t coins of this pile
        groups.append(items)
    return group_knapsack(groups, k)

When exactly one item per group must be chosen (for example "number of ways to roll dice to a target sum"), use a fresh array per group: new[j] = sum(old[j - w] for each item), starting from old[0] = 1.

Two-dimensional cost

Add one dimension per budget and loop all of them descending:

def two_cost_knapsack(items: list[tuple[int, int, int]], cap_a: int, cap_b: int) -> int:
    # items: (cost_a, cost_b, value); f[a][b] = best value within both budgets
    f = [[0] * (cap_b + 1) for _ in range(cap_a + 1)]
    for ca, cb, v in items:
        for a in range(cap_a, ca - 1, -1):
            for b in range(cap_b, cb - 1, -1):
                f[a][b] = max(f[a][b], f[a - ca][b - cb] + v)
    return f[cap_a][cap_b]

"Largest set of binary strings with at most m zeros and n ones" is this exact model with value 1 per string.

Reconstructing the chosen items

A single reused row forgets which decision produced each value. Keep the full 2D table and walk backward: if f[i][j] differs from f[i-1][j], item i had to be taken.

def knapsack_with_items(weights: list[int], values: list[int], cap: int) -> tuple[int, list[int]]:
    n = len(weights)
    # f[i][j] = best value using the first i items with weight <= j
    f = [[0] * (cap + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        w, v = weights[i - 1], values[i - 1]
        for j in range(cap + 1):
            f[i][j] = f[i - 1][j]
            if j >= w and f[i - 1][j - w] + v > f[i][j]:
                f[i][j] = f[i - 1][j - w] + v
    chosen, j = [], cap
    for i in range(n, 0, -1):
        if f[i][j] != f[i - 1][j]:  # the value changed, so item i - 1 was taken
            chosen.append(i - 1)
            j -= weights[i - 1]
    return f[n][cap], chosen[::-1]

If memory is tight, keep the 1D values and a separate took[i][j] boolean table (one bit per cell) instead.

Counting optimal solutions

Track two arrays over "weight exactly j": the best value f[j] and the number of subsets g[j] achieving it. When a candidate is strictly better, it replaces the count; when it ties, the counts add up.

def count_optimal(weights: list[int], values: list[int], cap: int) -> tuple[int, int]:
    NEG = float("-inf")
    f = [0] + [NEG] * cap  # f[j] = best value with weight exactly j
    g = [1] + [0] * cap  # g[j] = number of subsets achieving f[j]
    for w, v in zip(weights, values):
        for j in range(cap, w - 1, -1):
            cand = f[j - w] + v
            if cand > f[j]:
                f[j], g[j] = cand, g[j - w]
            elif cand == f[j]:
                g[j] += g[j - w]
    best = max(f)
    return best, sum(c for val, c in zip(f, g) if val == best)

Using "exactly j" matters: with "at most j" the same subset would be counted once for every capacity it fits in.

Undoing an item in a counting knapsack

Item order does not matter, so any item can be viewed as the last one added. If full[s] = f[s] + f[s - w], where f counts subsets without that item, then f can be recovered with an ascending pass. This answers "for each item, how many subsets of the others sum to T" in O(n * T) total instead of O(n^2 * T).

def counts_without_each(weights: list[int], target: int) -> list[int]:
    full = [1] + [0] * target  # full[s] = subsets of all items with sum s
    for w in weights:
        for s in range(target, w - 1, -1):
            full[s] += full[s - w]
    res = []
    for w in weights:  # weights are positive
        f = full[:]
        for s in range(w, target + 1):  # ascending: invert one 0/1 step
            f[s] -= f[s - w]
        res.append(f[target])
    return res

Other twists, briefly

• Mixed knapsack: some items once, some unlimited, some up to k times. Items are independent, so use the matching loop per item: descending, ascending, or binary splitting.
• Dependencies ("an accessory can be bought only together with its main item"): build a tree from the dependencies and merge children's knapsacks into the parent; this is tree knapsack, covered with tree DP.
• k-th best value: store, for every capacity, the top k values in decreasing order and merge the "skip" and "take" lists with two pointers, O(n * W * k).

Complexity

Variant	Time	Space
Group knapsack	O(W * total items)	O(W)
Two budgets	O(n * A * B)	O(A * B)
Reconstruction	O(n * W)	O(n * W)
Counting optimal	O(n * W)	O(W)
Undo for every item	O(n * W)	O(W)

Common pitfalls

• Group loop order: capacity must be outside the items of a group; otherwise the "one per group" rule breaks.
• Zero weights: the group loop and the undo trick both assume positive weights. A zero-weight item reads the value it is updating.
• Reconstruction from a single reused row: one row cannot be walked backward; keep the table or the choice bits.
• Counting with "at most" semantics: produces double counting; switch to "exactly" and sum at the end.
• Modulo in undo: with counts taken modulo p, subtract and normalize: (f[s] - f[s - w]) % p.

Interview variations

• Prefix of each pile (take coins from the top of stacks), one option per category under a budget: group knapsack.
• Dice or choices summing to a target: group counting with exactly one item per group.
• Strings limited by zeros and ones, tasks limited by time and money: two-dimensional cost.
• Return the actual subset (partition into two equal halves, show the items): reconstruction.
• Number of best plans: counting optimal solutions.