Search code examples
pythonpython-itertools

How do I find all possible ways I can fit an array to 4 slots with the possibility of the array having more/less than 4 numbers?


Let's say I have 4 buckets and an array of numbers like [1,2,3,4,5]

|  ||  ||  ||  |
|__||__||__||__|
 1   2   3   4
 5   2   3   4
 1   5   3   4     
 etc...

I can also have less than 4 numbers like [1, 2, 3]

|  ||  ||  ||  |
|__||__||__||__|
 1   2   3   
 1   3   2   
     1   2   3     
 etc...

How do I find all possible combinations of the numbers in buckets (like [1,2,3,4] if length is >= 4 or [1,2,None,None] if length is < 4)?


Solution

  • You can use a recursive generator function:

    def combos(nums, buckets, c = []):
      if len(c) == buckets:
        yield c
      else:
        if len(c) + len(nums) < buckets:
          yield from combos(nums, buckets, c+[None])
        for i, a in enumerate(nums):
          yield from combos(nums[:i]+nums[i+1:], buckets, c+[a])
    
    print(list(combos([1, 2, 3, 4, 5], 4)))
    print(list(combos([1, 2, 3], 4)))
    

    Output:

    [[1, 2, 3, 4], [1, 2, 3, 5], [1, 2, 4, 3], [1, 2, 4, 5], [1, 2, 5, 3], [1, 2, 5, 4], [1, 3, 2, 4], [1, 3, 2, 5], [1, 3, 4, 2], [1, 3, 4, 5], [1, 3, 5, 2], [1, 3, 5, 4], [1, 4, 2, 3], [1, 4, 2, 5], [1, 4, 3, 2], [1, 4, 3, 5], [1, 4, 5, 2], [1, 4, 5, 3], [1, 5, 2, 3], [1, 5, 2, 4], [1, 5, 3, 2], [1, 5, 3, 4], [1, 5, 4, 2], [1, 5, 4, 3], [2, 1, 3, 4], [2, 1, 3, 5], [2, 1, 4, 3], [2, 1, 4, 5], [2, 1, 5, 3], [2, 1, 5, 4], [2, 3, 1, 4], [2, 3, 1, 5], [2, 3, 4, 1], [2, 3, 4, 5], [2, 3, 5, 1], [2, 3, 5, 4], [2, 4, 1, 3], [2, 4, 1, 5], [2, 4, 3, 1], [2, 4, 3, 5], [2, 4, 5, 1], [2, 4, 5, 3], [2, 5, 1, 3], [2, 5, 1, 4], [2, 5, 3, 1], [2, 5, 3, 4], [2, 5, 4, 1], [2, 5, 4, 3], [3, 1, 2, 4], [3, 1, 2, 5], [3, 1, 4, 2], [3, 1, 4, 5], [3, 1, 5, 2], [3, 1, 5, 4], [3, 2, 1, 4], [3, 2, 1, 5], [3, 2, 4, 1], [3, 2, 4, 5], [3, 2, 5, 1], [3, 2, 5, 4], [3, 4, 1, 2], [3, 4, 1, 5], [3, 4, 2, 1], [3, 4, 2, 5], [3, 4, 5, 1], [3, 4, 5, 2], [3, 5, 1, 2], [3, 5, 1, 4], [3, 5, 2, 1], [3, 5, 2, 4], [3, 5, 4, 1], [3, 5, 4, 2], [4, 1, 2, 3], [4, 1, 2, 5], [4, 1, 3, 2], [4, 1, 3, 5], [4, 1, 5, 2], [4, 1, 5, 3], [4, 2, 1, 3], [4, 2, 1, 5], [4, 2, 3, 1], [4, 2, 3, 5], [4, 2, 5, 1], [4, 2, 5, 3], [4, 3, 1, 2], [4, 3, 1, 5], [4, 3, 2, 1], [4, 3, 2, 5], [4, 3, 5, 1], [4, 3, 5, 2], [4, 5, 1, 2], [4, 5, 1, 3], [4, 5, 2, 1], [4, 5, 2, 3], [4, 5, 3, 1], [4, 5, 3, 2], [5, 1, 2, 3], [5, 1, 2, 4], [5, 1, 3, 2], [5, 1, 3, 4], [5, 1, 4, 2], [5, 1, 4, 3], [5, 2, 1, 3], [5, 2, 1, 4], [5, 2, 3, 1], [5, 2, 3, 4], [5, 2, 4, 1], [5, 2, 4, 3], [5, 3, 1, 2], [5, 3, 1, 4], [5, 3, 2, 1], [5, 3, 2, 4], [5, 3, 4, 1], [5, 3, 4, 2], [5, 4, 1, 2], [5, 4, 1, 3], [5, 4, 2, 1], [5, 4, 2, 3], [5, 4, 3, 1], [5, 4, 3, 2]]
    [[None, 1, 2, 3], [None, 1, 3, 2], [None, 2, 1, 3], [None, 2, 3, 1], [None, 3, 1, 2], [None, 3, 2, 1], [1, None, 2, 3], [1, None, 3, 2], [1, 2, None, 3], [1, 2, 3, None], [1, 3, None, 2], [1, 3, 2, None], [2, None, 1, 3], [2, None, 3, 1], [2, 1, None, 3], [2, 1, 3, None], [2, 3, None, 1], [2, 3, 1, None], [3, None, 1, 2], [3, None, 2, 1], [3, 1, None, 2], [3, 1, 2, None], [3, 2, None, 1], [3, 2, 1, None]]
    

    At each recursive call, if a combination has not yet been formed, the code does two things:

    1. Checks if the specified number of buckets is greater than the input number list. If so, then the output is padded with None.
    2. The remaining number list is iterated over, and each iteration value is added to the running result, and the recursion proceeds.