Finding the number of ways to partition {1,2, ..., N} into p1 and p2 such that sum(p1) == sum(p2)

Question

I am trying to write a space and time efficient algorithm to calculate the number of ways to partition a set of integers {1, 2, ..., N} into two partitions such that the sums of integers in the two partitions are equal. I have started with a brute force approach that seems to give me correct results:

from itertools import combinations
def f(n):
    s = set([i+1 for i in range(n)])
    valid_partitions = []
    for i in range(1,n):
        combos = combinations(s,i)
        for c in combos:
            p1 = frozenset(c)
            p2 = set(range(1,n+1)) - p1
            p2 = frozenset(p2)
            if sum(p1) == sum(p2):
                valid_partitions.append(frozenset((p1,p2)))

    return len(set(valid_partitions))

Also, since I am checking every way to partition {1,2,...,N} into two non-empty sets, is it correct to say that the Big O time complexity of this approach is S(N,2) (Stirling Number of the Second Kind)? Also, how would I evaluate the Big O space complexity of this approach? Is there a better approach I can take?

Answer 1

Duplicated code/values

You have both:

s = set([i+1 for i in range(n)])

and

p2 = set(range(1,n+1)) - p1

This is bad because you have the same sets generated in 2 places with 2 different pieces of code:

• it is not efficient
• it is not easy to read and understand
• it is easy to get wrong if you ever need to update the set definition.

You should write something like:

        p2 = frozenset(s - p1)

Various hints

If sum(p1) == sum(p2), then sum(p1) == sum(p2) == sum(s) // 2 and sum(p1) == sum(s) // 2 is a sufficient condition. You could write something like:

def f(n):
    s = set([i+1 for i in range(n)])
    target_sum = sum(s) / 2
    valid_partitions = []
    for i in range(1,n):
        for c in combinations(s,i):
            if sum(c) == target_sum:
                p1 = frozenset(c)
                p2 = frozenset(s - p1)
                valid_partitions.append(frozenset((p1,p2)))
    return len(set(valid_partitions))

You could just consider only p1 without bothering about p2. Partitions would be counted twice but you can always divide the result at the end. Once you start this, you realise that all the logic about removing duplicates can be removed (all generation combinations are unique): you do not need sets and you do not need frozensets:

def f(n):
    s = set([i+1 for i in range(n)])
    target_sum = sum(s) / 2
    valid_partitions = []
    for i in range(1,n):
        for c in combinations(s,i):
            if sum(c) == target_sum:
                valid_partitions.append(c)
    return len(valid_partitions) // 2

Then, you do not need the have all partitions in a container, a simple counter will do the trick:

def f(n):
    s = set([i+1 for i in range(n)])
    target_sum = sum(s) / 2
    nb_partitions = 0
    for i in range(1,n):
        for c in combinations(s,i):
            if sum(c) == target_sum:
                nb_partitions += 1
    return nb_partitions // 2

Then, if you want to make things more concise, you could use the sum builtin and use the fact that True is 1 and False is 0 to write something like:

def f(n):
    s = set([i+1 for i in range(n)])
    target_sum = sum(s) / 2
    return sum(sum(c) == target_sum for i in range(1, n) for c in combinations(s, i)) // 2

Final optimisation

If you look at the value returned for the first 19 values, you get something like:

0 0
1 0
2 0
3 1
4 1
5 0
6 0
7 4
8 7
9 0
10 0
11 35
12 62
13 0
14 0
15 361
16 657
17 0
18 0
19 4110

We can see that the functions return 0 in 1 case out of 2. This corresponds to the fact that we cannot partition an odd sum. This can be taken into account to avoid computing combinations when this happens:

def f(n):
    s = set([i+1 for i in range(n)])
    target_sum, rem = divmod(sum(s), 2)
    if rem:
        return 0
    return sum(sum(c) == target_sum for i in range(1, n) for c in combinations(s, i)) // 2

Last detail

Your first call to set is useless because the elements are unique already and you just need an iterable. Also, you could get rid of the list comprehension by tweaking the parameters given to range: s = range(1, n+1).

Answer 2

Also, since I am checking every way to partition {1,2,...,N} into two non-empty sets, is it correct to say that the Big O time complexity of this approach is S(N,2) (Stirling Number of the Second Kind)?

No. Firstly, as pointed out by Josay, the algorithm considers each set partition twice. But secondly, it doesn't perform \$O(1)\$ work per set partition. It's performing \$\Theta(n)\$ work in calculating the sums. As far as I can see the documentation of itertools doesn't guarantee anything about how long it takes to generate each partition, but Gareth Rees points out that the implementation is also \$\Theta(n)\$. So the complexity is \$\Theta(n{n \brace 2}) = \Theta(n \, 2^n)\$.

Is there a better approach I can take?

Yes. The decision variant of this challenge is the subset sum problem and has a well-known dynamic programming solution. That solution is easily adapted to give counts rather than just Booleans. In pseudo-code:

counts = { 0 => 1 }
for elt in range(1, n+1):
    temp = { sum + elt => freq for (sum, freq) in counts }
    possibleSums = keys(counts).union(keys(temp))
    counts = { counts[sum] + temp[sum] for sum in possibleSums }
return counts[n * (n + 1) / 4]

In the worst case counts has \$2^\textrm{elt-1}\$ keys at the start of the loop body, so this runs in \$O(2^n)\$ time assuming fast dictionary lookups.

For completeness I will note that an alternative, but more complicated, solution in \$\Theta(2^n)\$ time would be to go through the subsets using a Gray code.

For even more efficient solutions to this particular case (where the set in question is range(1, n+1) rather than an arbitrary set), you could look at porting code from OEIS A058377.

Answer 3

is it correct to say that the Big O time complexity of this approach is S(N,2) (Stirling Number of the Second Kind)?

Since each element can be put in p1 or p2, there are two options for each element, giving a total of 2ⁿ different options. Since you are avoiding the cases where p1 or p2 are empty, it's 2ⁿ-2, which is twice the Stirling Number of the Second Kind (you're double counting since the Sterling number is based on p1 and p2 being indistinguishable). Since Big O ignores adding or subtracting constants, it's simpler to just give the complexity in terms of 2ⁿ rather than putting it in term of Sterling Numbers. And as @Peter Taylor points out, you are looking only at the number of iterations, and ignoring the complexity within an iteration.

You can avoid most of the double-counting by taking combinations up to n//2. Anything past that will be just a partition you've already gotten with p1 and p2 swapped. At i == n//2, you will be double-counting, so you can use @Peter Taylor's suggestion and set aside one element and then add it back in to a fixed partition.