ring_game.py

from itertools import permutations   
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

num_rings = 5

def find_config(num_rings):
    rings = [i for i in range(num_rings)]
    for i in range(num_rings-1):
        rings.append(7000)

    configs = []
    for p in permutations(rings,num_rings*2-1):
        configs.append(p)

    return configs

def check_configs(configs):

    v_configs = set()

    for c in configs:
        valid = check_config(c)
        if valid:
            c = c[num_rings-1:]
            v_configs.add(c)

    return v_configs

def check_config(config):
    valid = True
    stack = config[num_rings-1:]
    if not (stack[num_rings-1] == 7000 or stack[num_rings-1] == num_rings-1): valid = False
    for i in range(num_rings-1,0,-1):
        pos = i - 1
        ring = stack[i-1]
        ring_below = stack[i]
        if (ring != 7000) and (ring != pos) and (ring-1 != ring_below): valid = False

    return valid

def draw_config(configs):
    plt.figure(figsize=(10, 10))

    for fs in range(len(configs)):

        ax = plt.subplot(8,13,fs+1)

        config = configs[fs]

        width_bot = 0.4
        width_top = 0.1

        verts = [(0.5-width_bot/2,0),(0.5+width_bot/2,0),(0.5+width_top/2,1),(0.5-width_top/2,1)]
        poly = Polygon(verts,facecolor = '0.9',edgecolor='0.7')
        ax.add_patch(poly)

        colors = ['red','orange','green','blue','purple','pink','cyan']

        for c in range(len(config)):
            ring = config[c]
            pos = c

            if ring != 7000:
                height_b = 1 - ((pos + 1)/ len(config))
                height_t = 1 - (pos) / len(config)

                width_addition = 0.01


                width_l = 0.5 - ((ring+1)/len(config)) * (width_bot)/2 - (1-((ring+1)/len(config))) * (width_top)/2
                width_r = 0.5 + ((ring+1)/len(config)) * (width_bot)/2 + (1-((ring+1)/len(config))) * (width_top)/2

                width_l -= width_addition
                width_r += width_addition

                verts = [(width_l,height_b),(width_l,height_t),(width_r,height_t),(width_r,height_b)]
                poly = Polygon(verts,color = colors[ring], alpha = 0.5)
                ax.add_patch(poly)

        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_xlim([0,1])
        ax.set_ylim([0,1])

    plt.show()

configs = find_config(num_rings)
valid_configs = check_configs(configs)

print(len(valid_configs))

draw_config(list(valid_configs))