递归生成所有数位和为 n 的 k 位数字

Question

I was working on a problem where I'm finding all k-digit numbers whose digits sum up the given n.

I found how to do this and approach it as Integer Partitioning problem, however I would like to be able to only input n and k numbers (without the max_element) but when I try to delete it from the code it doesn't seem to work anymore.

How can I change that plus reverse it?

def c(n, k, max_element):
    allowed = range(max_element, 0, -1)

    def helper(n, k, t):
        if k == 0:
            if n == 0:
                yield t
        elif k == 1:
            if n in allowed:
                yield t + (n,)
        elif 1 * k <= n <= max_element * k:
            for v in allowed:
                yield from helper(n - v, k - 1, t + (v,))
    return helper(n, k, ())

for p in c(5, 3, 3):
    print(p)

I tried using the reversed method but apparently it doesn't work in the generator.

Result:

(3, 1, 1)
(2, 2, 1)
(2, 1, 2)
(1, 3, 1)
(1, 2, 2)
(1, 1, 3)

Expected result:

113 122 131 212 221 311

Answer 1

There are a couple of problems here; the first is that you want the numbers in order and this code generates them in reverse order, because of range(max_element, 0, -1) . The other problem is that since you're generating digits, the minimum element should be 0 and the maximum element should always be 9. We can fix both by changing that range to range(10) .

We still need to be careful not to generate numbers starting with 0, so we'll make allowed a parameter and use range(1, 10) for just the first digit. I've also changed it to return the result as an integer instead of a tuple.

For reference, the code for this generator function comes from my answer to another question .

def c(n, k):
    def helper(n, k, t, allowed):
        if k == 0:
            if n == 0:
                yield t
        elif k == 1:
            if n in allowed:
                yield 10*t + n
        elif 0 <= n <= 9 * k:
            for v in allowed:
                yield from helper(n - v, k - 1, 10*t + v, range(10))

    return helper(n, k, 0, range(1, 10))

Example:

>>> for p in c(5, 3):
...     print(p)
...
104
113
122
131
140
203
212
221
230
302
311
320
401
410
500

Answer 2

This function should do the trick

def c(n, k, max_element):
    allowed = range(max_element, 0, -1)

    def helper(n, k, t):
        if k == 0:
            if n == 0:
                yield t
        elif k == 1:
            if n in allowed:
                yield t + (n,)
        elif 1 * k <= n <= max_element * k:
            for v in allowed:
                yield from helper(n - v, k - 1, t + (v,))
    return helper(n, k, ())

def reversed_iterator(iter):
    return reversed(list(iter))

for p in reversed_iterator(c(5, 3, 3)):
    print(p)

here is the output :

(1, 1, 3)
(1, 2, 2)
(1, 3, 1)
(2, 1, 2)
(2, 2, 1)
(3, 1, 1)