Generate restricted weak integer compositions (or partitions) of an integer n into k parts in Python

后端 未结 1 1060
余生分开走
余生分开走 2020-12-21 18:36

(Re-posting, as I did not get any response to my previous post)
I am trying to write a Python code to generate weak integer compositions (partitions) of a number \'n\' i

相关标签:
1条回答
  • 2020-12-21 18:42

    This kind of problem is most straightforward to solve with a recursive generator function. To generate partitions of n into k parts, we can select the first part v, and then recursively partition n - v into k - 1 parts.

    You want earlier solutions to have larger numbers in the first position, so we'll choose v in descending order.

    def constrained_partitions(n, k, min_elem, max_elem):
        allowed = range(max_elem, min_elem-1, -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 min_elem * k <= n <= max_elem * k:
                for v in allowed:
                    yield from helper(n - v, k - 1, t + (v,))
    
        return helper(n, k, ())
    

    Example:

    >>> for p in constrained_partitions(5, 3, 0, 3):
    ...     print(p)
    ... 
    (3, 2, 0)
    (3, 1, 1)
    (3, 0, 2)
    (2, 3, 0)
    (2, 2, 1)
    (2, 1, 2)
    (2, 0, 3)
    (1, 3, 1)
    (1, 2, 2)
    (1, 1, 3)
    (0, 3, 2)
    (0, 2, 3)
    
    0 讨论(0)
提交回复
热议问题