play hue drops with matplotlib and test solver

huedrops.py

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np

colours = {'R':0.00,
           'O':0.14,
           'Y':0.22,
           'G':0.37,
           'B':0.77,
           'V':0.88}

##### random grid #######
def getCharFromNum(x):
    terv = 1./6
    if x<terv:
        return 'R'
    elif x<2*terv:
        return 'O'
    elif x<3*terv:
        return 'Y'
    elif x<4*terv:
        return 'G'
    elif x<5*terv:
        return 'B'
    else:
        return 'V'

def getRandomGrid(c, r):
    x = np.random.rand(c, r)
    y = np.chararray((c,r), unicode=True)
    for index, n in np.ndenumerate(x):
        y[index] = getCharFromNum(n)
    y[0,0] = 'W'
    return y

#### Challenge input #####
def getGridProblem():
    x = np.chararray((12,10), unicode=True)
    x[0,:] = ['W','Y','O','B','V','G','V','Y','O','B']
    x[1,:] = ['G','O','O','V','R','V','R','G','O','R']
    x[2,:] = ['V','B','R','R','R','B','R','B','G','Y']
    x[3,:] = ['B','O','Y','R','R','G','Y','V','O','V']
    x[4,:] = ['V','O','B','O','R','G','B','R','G','R']
    x[5,:] = ['B','O','G','Y','Y','G','O','V','R','V']
    x[6,:] = ['O','O','G','O','Y','R','O','V','G','G']
    x[7,:] = ['B','O','O','V','G','Y','V','B','Y','G']
    x[8,:] = ['R','B','G','V','O','R','Y','G','G','G']
    x[9,:] = ['Y','R','Y','B','R','O','V','O','B','V']
    x[10,:] = ['O','B','O','B','Y','O','Y','V','B','O']
    x[11,:] = ['V','R','R','G','V','V','G','V','V','G']
    return x

### Slow fill mechanics ####
def getAdjacent(x, ind, direction):
    if direction==0 and ind[0]<x.shape[0]-1:
        #down
        return (ind[0]+1,ind[1])
    elif direction==1 and ind[1]<x.shape[1]-1:
        #right
        return (ind[0],ind[1]+1)
    elif direction==2 and ind[0]>0:
        #up
        return (ind[0]-1,ind[1])
    elif direction==3 and ind[1]>0:
        #left
        return (ind[0],ind[1]-1)
    else:
        return None

def getMatches(x):
    ind = [(0,0)]
    found = False
    while True:
        for direction in range(4):
            for i in ind:
                adj = getAdjacent(x, i, direction)
                if adj!=None:
                    if x[adj]==x[i]:
                        if adj not in ind:
                            ind.append(adj)
                            found = True
        if found:
            found = False
        else:
            break
    return ind

def fill(x, c):
    newx = x.copy()
    for index in getMatches(x):
        newx[index] = c
    return newx

#### greedy solver, thinks two moves ahead #####
def solver(x):
    most = 0
    route = ('W','W')
    routes = []
    newx = x.copy()
    bestx = x.copy()
    while len(getMatches(newx))<x.shape[0]*x.shape[1]:
        for colour1 in colours.keys():
            x1 = newx.copy()
            x1 = fill(x1, colour1)
            for colour2 in colours.keys():
                x2 = x1.copy()
                x2 = fill(x2, colour2)
                if len(getMatches(x2))>most:
                    most = len(getMatches(x2))
                    route = colour1, colour2
                    bestx = x2.copy()
        routes.append(route)
        newx = bestx.copy()
    print(routes)
    
#### matplotlib display #####
def getColourNum(c):
    try:
        return colours[c]
    except KeyError:
        return 1

def getColourGrid(x):
    y = np.zeros(x.shape)
    for index, n in np.ndenumerate(x):
        y[index] = getColourNum(n)
    return y

def plot():
    ax.clear()
    ax.imshow(getColourGrid(X), cmap=plt.get_cmap('gist_rainbow'), interpolation='nearest', norm=nrm)
    ax.set_axis_off()
    plt.show()

def onpress(event):
    global X
    for colour in colours.keys():
        if event.key == colour.lower() or event.key == colour:
            X = fill(X, colour)
    plot()

#### run it #####

#X = getRandomGrid(5,4)
X = getGridProblem()
solver(X)
mpl.rcParams['toolbar'] = 'None'
plt.rcParams["figure.figsize"] = [X.shape[1]*0.8,X.shape[0]*0.8]
nrm = colors.NoNorm()
fig = plt.figure()
ax = plt.Axes(fig, [0,0,1,1])
fig.add_axes(ax)
fig.canvas.mpl_connect('key_press_event', onpress)
plot()