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)
.