Topics › Dynamic Programming › DP on Subsequences

Count partitions with given difference

Hard TUF ID 317 A2Z

Problem

You are given an array arr of n non-negative integers and a non-negative integer d. Put every element into exactly one of two groups S1 and S2 so that sum(S1) ≥ sum(S2) and sum(S1) - sum(S2) = d. Count the number of such partitions, modulo 10^9 + 7. Two partitions are different when some index lands in a different group, so equal values at different positions give different partitions.

Typical constraints: 1 ≤ n ≤ 50, 0 ≤ arr[i] ≤ 50, 0 ≤ d ≤ 2500. Zeros are allowed.

Example 1

arr = [5, 2, 6, 4], d = 3
answer = 1          // S1 = {6, 4} = 10, S2 = {5, 2} = 7

Example 2

arr = [1, 1, 1, 1], d = 0
answer = 6          // pick which two of the four 1s go into S2: C(4, 2) = 6

Approach 1: Recursion over every assignment

Intuition. The definition gives a direct brute force: every element goes to S1 or to S2, so there are 2^n assignments. Walk through the elements, carry the running sum s1 of S1, and at the end check whether s1 - (total - s1) equals d. Since d ≥ 0, that equality already implies sum(S1) ≥ sum(S2).

Algorithm.

  1. Compute total.
  2. assign(i, s1): if i == n, return 1 when 2 · s1 - total == d, else 0.
  3. Otherwise return assign(i + 1, s1 + arr[i]) + assign(i + 1, s1), modulo MOD.
static int recursive(int[] arr, int d) {
    int total = 0;
    for (int x : arr) total += x;
    return assign(0, 0, arr, total, d);
}
static int assign(int i, int s1, int[] arr, int total, int d) {
    if (i == arr.length) return (s1 - (total - s1) == d) ? 1 : 0;
    int toS1 = assign(i + 1, s1 + arr[i], arr, total, d);
    int toS2 = assign(i + 1, s1, arr, total, d);
    return (toS1 + toS2) % MOD;
}

Complexity. Time O(2^n), space O(n). Adding a memo on (i, s1) would already make this polynomial, but there is a better formulation that needs a smaller table and reuses a problem we already know.

The reduction to counting subsets

Write S1 and S2 for the two sums. We know two equations:

S1 + S2 = total
S1 - S2 = d

Subtracting the second from the first gives 2 · S2 = total - d, so

S2 = (total - d) / 2

A partition is completely determined by which elements form S2 (everything else is S1). Therefore: the number of partitions with difference d equals the number of subsets with sum (total - d) / 2. The problem becomes Count subsets with sum K with K = (total - d) / 2.

Two quick rejections come out of the formula:

smallerSum performs both checks and returns -1 when either fails.

static int smallerSum(int[] arr, int d) {
    int total = 0;
    for (int x : arr) total += x;
    if (d > total || (total - d) % 2 != 0) return -1;
    return (total - d) / 2;
}

Why target S2 rather than S1 = (total + d) / 2? Both give the same count, because a subset with sum S2 and its complement with sum S1 are in one-to-one correspondence. But S2 ≤ total / 2, so the DP table is at most half as wide.

Approach 2: Memoization

Intuition. Count subsets of the first i elements with sum t, by skipping or taking arr[i - 1]. The two branches never produce the same subset, so their counts add. The base case is on the prefix length: with i = 0 elements, only the empty subset exists, and it has sum 0.

That base case matters because zeros are allowed. A shortcut like “if t == 0 return 1” would stop early and ignore that each remaining zero could join S2 or stay in S1, two genuinely different partitions. For [0, 0, 1] and d = 1 the correct answer is 4, and the shortcut would report 1.

Algorithm.

  1. target = smallerSum(arr, d); return 0 if it is -1.
  2. Allocate dp[n + 1][target + 1] filled with -1.
  3. memo(i, t): base case i == 0; otherwise cached notTake + take, modulo MOD.
static int memoization(int[] arr, int d) {
    int target = smallerSum(arr, d);
    if (target < 0) return 0;
    int[][] dp = new int[arr.length + 1][target + 1];
    for (int[] row : dp) Arrays.fill(row, -1);
    return memo(arr.length, target, arr, dp);
}
static int memo(int i, int t, int[] arr, int[][] dp) {
    if (i == 0) return t == 0 ? 1 : 0;              // prefix-length base case: zeros are counted correctly
    if (dp[i][t] != -1) return dp[i][t];
    int notTake = memo(i - 1, t, arr, dp);
    int take = arr[i - 1] <= t ? memo(i - 1, t - arr[i - 1], arr, dp) : 0;
    return dp[i][t] = (notTake + take) % MOD;
}

Complexity. Time O(n · target), space O(n · target) plus O(n) stack, where target = (total - d) / 2 ≤ total / 2.

Approach 3: Tabulation

Intuition. Fill dp[i][t] row by row from dp[0][0] = 1. The inner loop starts at t = 0, because a zero element creates new subsets with sum 0 and so column 0 can grow.

Algorithm.

  1. Reject impossible inputs via smallerSum.
  2. dp[0][0] = 1; for i = 1..n, t = 0..target: dp[i][t] = dp[i - 1][t] + (arr[i - 1] ≤ t ? dp[i - 1][t - arr[i - 1]] : 0), modulo MOD.
  3. Return dp[n][target].
static int tabulation(int[] arr, int d) {
    int target = smallerSum(arr, d);
    if (target < 0) return 0;
    int n = arr.length;
    int[][] dp = new int[n + 1][target + 1];
    dp[0][0] = 1;
    for (int i = 1; i <= n; i++) {
        for (int t = 0; t <= target; t++) {
            int notTake = dp[i - 1][t];
            int take = arr[i - 1] <= t ? dp[i - 1][t - arr[i - 1]] : 0;
            dp[i][t] = (notTake + take) % MOD;
        }
    }
    return dp[n][target];
}

Dry run on Example 1: total = 17, d = 3, so target = 7. Only the columns that become non-zero are shown.

i (element) t = 0 2 4 5 6 7
0 (none) 1 0 0 0 0 0
1 (5) 1 0 0 1 0 0
2 (2) 1 1 0 1 0 1
3 (6) 1 1 0 1 1 1
4 (4) 1 1 1 1 2 1

dp[4][7] = 1: the single subset {5, 2} forms S2, matching the example.

Complexity. Time O(n · target), space O(n · target).

Approach 4: Space optimized

Intuition. Keep only one row, ways. For each element x, sweep t from target down to x and add ways[t - x] into ways[t]. The downward direction guarantees that ways[t - x] has not yet been updated for x, so x is used at most once. For x = 0 the sweep doubles every entry, which is the “zero on either side” effect.

Algorithm.

  1. Reject impossible inputs.
  2. ways[0] = 1; for each x, for t from target down to x: ways[t] = (ways[t] + ways[t - x]) % MOD.
  3. Return ways[target].
static int spaceOptimized(int[] arr, int d) {
    int target = smallerSum(arr, d);
    if (target < 0) return 0;
    int[] ways = new int[target + 1];
    ways[0] = 1;
    for (int x : arr) {
        for (int t = target; t >= x; t--) ways[t] = (ways[t] + ways[t - x]) % MOD;
    }
    return ways[target];
}

Complexity. Time O(n · target), space O(target).

Edge cases and pitfalls