Counting Colours (AMO 2022 P4).py

import random

# generate all 2^n possible combos for the new row
def make_possible_rows(existing_rows, n):
  new_rows = []
  for i in range(len(existing_rows)):
    new_rows += [existing_rows[i] + ['B']]
    new_rows += [existing_rows[i] + ['R']]
  if len(new_rows) == 2**n:
    return new_rows
  else:
    return make_possible_rows(new_rows, n)
  
# check if a grid meets the conditions for colouring
def check_grid(grid):
  k = len(grid[-1])
  for j in range(1,k):
    for i in range(0,j):
      # (x,y) coords are reversed because grid is an array of rows
      if grid[j][i] == grid[k][j] and grid[j][i] != grid[k][i]:
        return False
  else:
    return True

def make_colourings(existing_colourings, n):
  new_colourings = []
  num_cols = len(existing_colourings[0])
  possible_rows = make_possible_rows([['B'],['R']], num_cols)
  for grid in existing_colourings:
    for new_row in possible_rows:
      if check_grid(grid + [new_row]):
        #print('this grid is OK:')
        #print(grid + [new_row])
        new_colourings += [grid + [new_row]]
      else:
        #print('this grid is NOT OK:')
        #print(grid + [new_row])
        continue
  if len(new_colourings[0]) == n:
    print("Done! The number of colourings where i,j<=" + str(n) + " is " + str(len(new_colourings)) + "\n")
    print("Here is one of the colourings...")
    example = random.choice(new_colourings)
    example.reverse()
    print(*example,sep='\n')
  else:
    return make_colourings(new_colourings, n)

make_colourings([ [[' '],['B']], [[' '],['R']]  ],6)
# OUTPUT: 
#
# Done! The number of colourings where i,j<=6 is 720
#
# Here is one of the colourings...
# ['R', 'R', 'R', 'R', 'R']
# ['B', 'R', 'R', 'B']
# ['R', 'R', 'R']
# ['B', 'B']
# ['B']
# [' ']