Since the values apparently are evenly spaced and start at 0, this is equivalent to a problem in integers: we want to find three ordered values from range(6) which sum to 5. Aside from avoiding floating-point inaccuracy, this offers a significant performance improvement in my testing (even without considering the overhead of math.isclose). I will show code for this problem, and trust that it can be adapted as needed.
For reference, integer implementations of the Numpy and itertools approaches:
import numpy as np
from itertools import combinations_with_replacement
def convex_itertools(count, total):
candidates = combinations_with_replacement(range(total + 1), r=count)
return [c for c in candidates if sum(c) == total]
def convex_numpy(count, total):
mesh = np.meshgrid(*([range(total + 1)] * count))
candidates = np.stack(mesh, -1).reshape(-1, count)
return candidates[np.sum(candidates, axis=1) == total]
On my machine:
>>> from timeit import timeit
>>> timeit('convex_itertools(3, 5)', globals=globals(), number=10000)
0.08309134596493095
>>> timeit('convex_numpy(3, 5)', globals=globals(), number=10000)
0.8170730160200037
The primary problem here is algorithmic. Using meshgrid produces unordered results, and as such produces many more results than desired. But even using itertools.combinations_with_replacement, there is a lot of filtering to do. For the sample problem size, already 56 ordered candidates (216 unordered) are generated, which will filter down to only 5 (21, respectively).
Thus, it would be preferable to generate the desired results directly, rather than via filtering. For this, I will use a recursive generator:
def convex_manual(count, total):
if count == 1:
if total >= 0:
yield [total]
return
for first in range((total // count) + 1):
remainder = total - (first * count)
for result in convex_manual(count - 1, remainder):
yield [first] + [first+r for r in result]
Our base case is when the combinations consist of a single value: there is either a single solution (just use the full total, when it's non-negative and thus a legal value), or no solutions (previously chosen values in the recursion already add up to more than the total desired amount - though this should be impossible from how the recursion is constructed, except with bad initial values). For combinations of more than one value, we consider the possible first values to use (it cannot exceed the average value), recursively find the ways to distribute the excess within our limit (after accounting for the fact that each number must be at least as large as the first, or else the output will not be ordered), and translate those recursive findings into results for the current step.
Let's test it:
>>> list(convex_manual(3, 5))
[[0, 0, 5], [0, 1, 4], [0, 2, 3], [1, 1, 3], [1, 2, 2]]
>>> list(convex_manual(3, 10))
[[0, 0, 10], [0, 1, 9], [0, 2, 8], [0, 3, 7], [0, 4, 6], [0, 5, 5], [1, 1, 8], [1, 2, 7], [1, 3, 6], [1, 4, 5], [2, 2, 6], [2, 3, 5], [2, 4, 4], [3, 3, 4]]
Even for the small input example, and despite the overhead of assembling many lists and using a recursive generator, this is faster:
>>> timeit('list(convex_manual(3, 5))', globals=globals(), number=10000)
0.07173079502535984
Of course, considering how the recursion works gives some insight into improvements for the other approaches. In particular: there is only ever one possibility for the last value, so we should only generate candidates for the other values and then see if they lead to a legal combination. Thus:
import numpy as np
from itertools import combinations_with_replacement
def convex_itertools_opt(count, total):
candidates = combinations_with_replacement(range(total + 1), r=count-1)
return [
c + (remainder,) for c in candidates
if (remainder := total - sum(c)) >= c[-1]
]
def convex_numpy_opt(count, total):
mesh = np.meshgrid(*([range(total + 1)] * (count-1)))
candidates = np.stack(mesh, -1).reshape(-1, (count-1))
# Reference for the technique here:
# https://stackoverflow.com/questions/8486294
candidates = np.c_[candidates, total - np.sum(candidates, axis=1)]
return candidates[candidates[:, -1] >= 0]
Which does get better results for the itertools approach:
>>> timeit('convex_itertools_opt(3, 5)', globals=globals(), number=10000)
0.053549989010207355
but of course will break down for larger problem sizes:
>>> timeit('convex_itertools_opt(5, 10)', globals=globals(), number=10000)
1.7618601340218447
>>> timeit('list(convex_manual(5, 10))', globals=globals(), number=10000)
0.6665365709923208
and incurs too much overhead in Numpy:
>>> timeit('convex_numpy_opt(3, 5)', globals=globals(), number=10000)
0.9056216909666546
Contrary to what one might expect, trying to use lists and list comprehensions at every step (rather than using a recursive generator and explicitly creating a list from the top-level result) seems to have little effect on performance, regardless of the problem size. I tried using Numpy tools for gathering the results and performance got much, much worse and the code was also much harder to understand.
