Search code examples
pythonnumpypython-itertools

Coloring n boxes with m colors


I have n boxes I want to color them with m colors. I want to allow repeat of colors. For example given 4 boxes and two colors. Denoting the colors by 1 and 2 we have the following ways to color them

[[1, 1, 1, 1], [1, 1, 1, 2], [1, 1, 2, 1], [1, 1, 2, 2], [1, 2, 1, 1], 
[1, 2, 1, 2], [1, 2, 2, 1], [1, 2, 2, 2], [2, 1, 1, 1], [2, 1, 1, 2],
[2, 1, 2, 1], [2, 1, 2, 2], [2, 2, 1, 1], [2, 2, 1, 2], [2, 2, 2, 1],
[2, 2, 2, 2]]

where for instance [1,1,1,1] means coloring box 1 with first color and box 2 with first color up to the last box. While [1, 1, 2, 1] means coloring boxes 1,2 and 4 with color 1 while box 3 with color 2.

To that end I defined the following function

def recursive_fun(number_of_boxes,number_of_colors):
      possible_colors=range(1,number_of_colors+1)
      if number_of_boxes==1:
          return [[i] for i in possible_colors]
      else:
          output=[] 
          y=recursive_fun(number_of_boxes-1,number_of_colors)
          for i in y:
               for m in possible_colors:
                     output.append(i+[m])
      return output

The function is working but I would like to have a more efficient way of doing this. Is there a way of doing this using the itertools package?


Solution

  • You mean like itertools.product?

    import itertools
    
    colours = (1, 2)
    
    for x in itertools.product(colours, repeat=4):
        print(x)
    

    prints:

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