How to generate all possible variations of a 2 matrix of an enum element in python?

sg_rs 0
Viewed 23

I have an enum:

class MyEnum(enum.Enum):
   EMTPY = 0
   APPLE = 1
   HONEY = 2

I want to feed in 2 values height and width and I want to get every possible variation of a 2D array where each element is a MyEnum.

E.g.

def variations(height, width):
    variation = # code
    yield variation

for v in variations(2,2):
    print(v)
    print("---------")

will print

[
  [EMTPY, EMTPY],
  [EMTPY, EMTPY]
]
---------
[
  [APPLE, EMTPY],
  [EMTPY, EMTPY]
]
---------
[
  [HONEY, EMTPY],
  [EMTPY, EMTPY]
]
---------
[
  [EMTPY, APPLE],
  [EMTPY, EMTPY]
]
...

I was looking into combinations and permutations of the itertools library, but every attempt failed, because this task seems more complex than the tutorials in the internet.

Because I will probably have big height and width parameters, I want my variations function to be a generator.

itertools combinations permutation enums python
1 Answers

You could make it in two steps:

  1. Use itertools.product() to generate width*height-sized permutations with repetition
  2. Split the list into rows of a matrix.
import itertools
def variations(items,width,height):
  for one in itertools.product(items,repeat=width*height):                # step 1
    yield [list(one[i*width:(i+1)*width]) for i in range(height)]         # step 2

for mtx in variations(["EMPTY","APPLE","HONEY"],2,2):
  print(mtx)

Will produce all the 81 (34) possible matrices:

[['EMPTY', 'EMPTY'], ['EMPTY', 'EMPTY']]
[['EMPTY', 'EMPTY'], ['EMPTY', 'APPLE']]
[['EMPTY', 'EMPTY'], ['EMPTY', 'HONEY']]
[['EMPTY', 'EMPTY'], ['APPLE', 'EMPTY']]
...
[['HONEY', 'HONEY'], ['HONEY', 'APPLE']]
[['HONEY', 'HONEY'], ['HONEY', 'HONEY']]
tevemadar0
Related