KenKen.py

# KenKen/Inky Solver using NumberJack - Jim Bumgardner

from Numberjack import *

#
# Ensure that the sum of the segments
# in cc == res
#

bStats = {'cnt':0,'failures':0,'backtracks':0,'walltime':0,'propags':0,'checks':0}

def parsePuzzle(sz,puzz):
  (layout,clues) = puzz.split(';')
  clues = clues.split(',')
  ldict = {}
  grid = []
  nbrClues = 0
  for i in range(sz*sz):
    ch = layout[i]
    if ch not in ldict:
      clue = clues[nbrClues]
      op = clue[-1]
      val = int(clue[:-1])
      cage = [val,op,[]]
      nbrClues += 1
      ldict[ch] = cage
      grid.append(cage)
    y = (i/sz)+1
    x = (i%sz)+1
    ldict[ch][2].append([y,x])
  return grid

  # problem = [
  #     [11, '+', [[1, 1], [2, 1]]],
  #     [2, '/',  [[1, 2], [1, 3]]],
  #     [20, '*', [[1, 4], [2, 4]]],
  #     [6, '-',  [[1, 5], [1, 6], [2, 6], [3, 6]]],
  #     etc...


def solveInky(sz,puzz, param):

  problem = parsePuzzle(sz,puzz)

  num_p = len(problem)

  grid = Matrix(sz, sz, 1, sz)
  model = Model()

  model += [AllDiff(row) for row in grid.row]
  model += [AllDiff(col) for col in grid.col]

  # calculate the segments
  for (res, op, segment) in problem:
    cage = [grid[y-1,x-1] for (y,x) in segment]
    if op == '+':
        model += Sum(cage) == res
    elif op == '-':
        model += ((cage[0] - cage[1]) == res) | ((cage[1] - cage[0]) == res)
    elif op == '*':
      # Note: Ideally, I'd prefer to do this:
      # model += Product(cage) == res
      ls = len(cage)
      if ls == 2:
        model += (cage[0] * cage[1]) == res
      elif ls == 3:
        model += (cage[0] * cage[1] * cage[2]) == res
      elif ls == 4:
        model += (cage[0] * cage[1] * cage[2]*cage[3]) == res
      elif ls == 5:
        model += (cage[0] * cage[1] * cage[2]*cage[3]*cage[4]) == res
      elif ls == 6:
        model += (cage[0] * cage[1] * cage[2]*cage[3]*cage[4]*cage[5]) == res
      else:
        print "Segment too large,len=",ls
      pass
    elif op == '/':
        model += ((cage[0] * res) == cage[1]) | ((cage[1] * res) == cage[0])
        # pass
    else:
      print "Invalid op",op

  solver = model.load(param['solver'])
  solver.reset(full=True) # needed to clear from previous solves
  solver.setVerbosity(param['verbose'])

  res = solver.solve()

  outS = ''
  # if solver.is_sat():
  if res == SAT:
    # return answer in string form
      outS = ''.join(map(str, grid.flat))  # shorter version of the following array:
      res2 = solver.getNextSolution()
      if res2 == SAT:
        return "mult"
      pass
  else:
      return "Unsatisfiable"

  bStats['cnt'] += 1
  bStats['failures'] += solver.getFailures()
  bStats['backtracks'] += solver.getBacktracks()
  bStats['propags'] += solver.getPropags()
  bStats['checks'] += solver.getChecks()
  bStats['walltime'] += solver.getTime()

  return outS


test_puzzles = [
(3,"AABCDBCDD;1-,3*,3*,7+","213321132"), # easy
(3,"AABCDBCDD;3*,1-,6*,6+","132321213"), # hard
(6,"ABCCDEABFGDEHHFGIIJJKLLMNNKOMMNNNOPP;2/,3-,6*,6+,2/,1-,18*,6+,3/,3/,7+,2-,7+,540*,4*,1-","413256245613154362621534536421362145"),   # easy
(6,"AABBCDEAFGCDEHFGIIJHKLMMJNKLOPNNQQOP;8+,20*,2/,7+,1-,3/,2/,2-,7+,1-,12*,3/,5+,6*,4*,1-,1-","124536453261361452546123612345235614"), # med
(6,"AAABCDEFBBCDEFGBBDHIGJKKHILJMNOOLMMN;72*,144*,2/,9+,3-,2-,1-,3/,7+,1-,5+,3/,40*,2-,6*","436125562413245361154632321546613254"),     # hard
]

if __name__ == "__main__":
  import time

  default = {'solver': 'Mistral', 'verbose': 0}
  param = input(default)

  st = time.time()

  nbrCorrect = 0
  nbrPuzzles = 0

  for (sz,puzz,ans) in test_puzzles:
      # print sz,puzz,ans
      lastAnswer = solveInky(sz,puzz,param)
      nbrPuzzles += 1
      if lastAnswer == ans:
        nbrCorrect += 1
        print "correct:",lastAnswer
      else:
        print "incorrect:",lastAnswer

  print "%d/%d %.2f secs" % (nbrCorrect, nbrPuzzles, time.time()-st),bStats