Find all combination that sum to N when multiplies by index
<p>Given a number of items (<em>n</em>), what is the most efficient way to generate all possible lists [a<sub>1</sub>, a<sub>2</sub>, ..., a<sub>n</sub>] of non-negative integers under the condition that:</p>

<p>1*a<sub>1</sub> + 2*a<sub>2</sub> + 3*a<sub>3</sub> + ... + n*a<sub>n</sub> = n </p>

<p>using Python?</p>

<p>So for example, given an n of 5, the following combinations are:</p>

<p>[0,0,0,0,1]</p>

<p>[1,0,0,1,0]</p>

<p>[0,1,1,0,0]</p>

<p>[2,0,1,0,0]</p>

<p>[1,2,0,0,0]</p>

<p>[3,1,0,0,0]</p>

<p>[5,0,0,0,0]</p>

<p>I've implemented a brute-force method that generates all permutations and then checks if the list meets the above requirement, but is there a more efficient way to do this?</p>

<p>A "greedy" algorithm works well for this.  I'm using Python 3 here:</p>

<pre><code>def pick(total):

    def inner(highest, total):
        if total == 0:
            yield result
            return
        if highest == 1:
            result[0] = total
            yield result
            result[0] = 0
            return
        for i in reversed(range(total // highest + 1)):
            result[highest - 1] = i
            newtotal = total - i * highest
            yield from inner(min(highest - 1, newtotal),
                             newtotal)

    result = [0] * total
    yield from inner(total, total)
</code></pre>

<p>Then, e.g.,</p>

<pre><code>for x in pick(5):
    print(x)
</code></pre>

<p>displays:</p>

<pre><code>[0, 0, 0, 0, 1]
[1, 0, 0, 1, 0]
[0, 1, 1, 0, 0]
[2, 0, 1, 0, 0]
[1, 2, 0, 0, 0]
[3, 1, 0, 0, 0]
[5, 0, 0, 0, 0]
</code></pre>

<p>Like most recursive algorithms, it does a more-or-less obvious thing, then recurses to solve the (sub)problem that remains.</p>

<p>Here <code>inner(highest, total)</code> means to find all the decompositions of <code>total</code> using integers no larger than <code>highest</code>.  How many copies of <code>highest</code> can we use?  The more-than-less obvious answer is that we can use 0, 1, 2, ..., up to (and including) <code>total // highest</code> copies, but no more than that.  Unless <code>highest</code> is 1 - then we have to use exactly <code>total</code> copies of 1.</p>

<p>However many copies we use of <code>highest</code>, the subproblem remaining is to decompose whatever remains of the total using integers no larger than <code>highest - 1</code>.  Passing <code>min(highest - 1, newtotal)</code> instead of <code>highest - 1</code> is an optimization, since it's pointless trying any integer larger than the new total.</p>

python algorithm permutation combinatorics