1 | initial version |
Here is an iterator over 4-tuples (a, b, c, d)
such that a + b + c + d = k
and a < b < c < d
. You can also ensure that d <= nmax
, which seems a requirement of your problem.
def get_series(k, nmax=None):
"""
Iterator over all series (a, b, c, d) such that:
- a + b + c + d = k
- a < b < c < d
When nmax is specified, we furthermore ensure that a,b,c,d <= nmax.
"""
if nmax is None:
for a in range(1, (k - 6) // 4 + 1):
for b in range(a + 1, (k - a - 3) // 3 + 1):
for c in range(b + 1, (k - a - b - 1) // 2 + 1):
yield (a, b, c, k - a - b - c)
else:
for a in range(1, min(nmax, (k - 6) // 4) + 1):
for b in range(a + 1, min(nmax, (k - a - 3) // 3) + 1):
for c in range(b + 1, min(nmax, (k - a - b - 1) // 2) + 1):
d = k - a - b - c
if d <= nmax:
yield (a, b, c, d)
you can check that
This is way faster than your brute force attempt
sage: %time NList = [i for i in Subsets([1..100], 4) if sum(i) == 100]
CPU times: user 31 s, sys: 83.3 ms, total: 31.1 s
Wall time: 31.1 s
sage: %time L = list(get_series(100))
CPU times: user 7.23 ms, sys: 136 µs, total: 7.36 ms
Wall time: 7.35 ms
sage: len(L) == len(NList)
True
sage: L == [tuple(sorted(i)) for i in NList]
True
This is of course not sufficient for solving efficiently your problem since you have plenty of 4-tuples to consider and then triples of 4-tuples...
sage: %time L = list(get_series(1000))
CPU times: user 2.82 s, sys: 99.1 ms, total: 2.92 s
Wall time: 2.92 s
sage: len(L)
6840777
Another useful idea is to store the tuples in a dictionary with key a^2 + b^2 + c^2 + d^2
and value the list of corresponding tuples (so with same sum and same sum of squares).