Solving Alphametics with z3

z3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving Alphametics with z3\n",
    "\n",
    "I came across [Alphametics](https://www.math.uni-bielefeld.de/~sillke/PUZZLES/ALPHAMETIC/alphametic-mike-keith.html), or [Verbal arithmetic](https://www.wikiwand.com/en/Verbal_arithmetic), quite long ago as I was looking for inspiration to write an assignement for my programming classes. The idea sat in the back of my head until recently I encountered the [Z3 theorem prover](https://github.com/Z3Prover/z3) while trying to solve [Day 24](https://adventofcode.com/2021/day/24) puzzle of [Advent of Code](https://adventofcode.com/). This short notebook put the two things together. It's a mere exercise I found fun enough to explore Z3 a bit more.\n",
    "\n",
    "An alphametics puzzle is an \"alphabetic equation\" such as, for example, `SEND + MORE = MONEY`; you are required to assign (distinct) decimal digits to letters so that the equation holds. A way to solve it is to translate it to a set of constraints on `S`, `E`, `N`, `D`, `M`, `O`, `R`, `Y` integer valued variables. A solver like Z3 can then be used to find an assignement (of integers to variables) that satisfies the constraints and hence solves the puzzle.\n",
    "\n",
    "Let's start doing some impots that will be used in the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from string import ascii_letters\n",
    "from z3 import Int, Solver, Distinct, sat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, we want to consider equations with more than two operands on the left side, using both additions and subtractions; to \"parse\" such expressions we'll use some Python standard library functions, here we wet our hands just splitting it in parts that will be processed later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEVEN\n",
      "+ SEVEN\n",
      "+ SEVEN\n",
      "+ NINE\n",
      "THIRTY\n"
     ]
    }
   ],
   "source": [
    "expr = 'SEVEN + SEVEN + SEVEN + NINE = THIRTY'\n",
    "\n",
    "left, right = expr.split('=')\n",
    "first, *rest = left.split() # we assume that parts are surrounded by spaces\n",
    "\n",
    "right = right.strip()\n",
    "\n",
    "print(first)\n",
    "for op, word in zip(rest[::2], rest[1::2]):\n",
    "  print(op, word)\n",
    "print(right)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A first thing to do is to create a variable for every letter; we keep such variables in a dictionary (mapping the letter name to the variabile itself):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "letters = set(expr) & set(ascii_letters)\n",
    "letter2variable = {letter: Int(letter) for letter in letters}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to translate a word in an expression representing its value (in terms of the above variables):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def word2expression(word):\n",
    "  return sum(10 ** pow * letter2variable[letter] for pow, letter in enumerate(word[::-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, the expression representing the word `SEND` is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "0 + 1·N + 10·E + 100·V + 1000·E + 10000·S"
      ],
      "text/plain": [
       "0 + 1*N + 10*E + 100*V + 1000*E + 10000*S"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "word2expression('SEVEN')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using such function (and the \"parse\" code above), we are ready to translate the \"alphabetic equation\" to an equation in our variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "0 +\n",
       "1·N +\n",
       "10·E +\n",
       "100·V +\n",
       "1000·E +\n",
       "10000·S +\n",
       "0 +\n",
       "1·N +\n",
       "10·E +\n",
       "100·V +\n",
       "1000·E +\n",
       "10000·S +\n",
       "0 +\n",
       "1·N +\n",
       "10·E +\n",
       "100·V +\n",
       "1000·E +\n",
       "10000·S +\n",
       "0 +\n",
       "1·E +\n",
       "10·N +\n",
       "100·I +\n",
       "1000·N =\n",
       "0 + 1·Y + 10·T + 100·R + 1000·I + 10000·H + 100000·T"
      ],
      "text/plain": [
       "0 +\n",
       "1*N +\n",
       "10*E +\n",
       "100*V +\n",
       "1000*E +\n",
       "10000*S +\n",
       "0 +\n",
       "1*N +\n",
       "10*E +\n",
       "100*V +\n",
       "1000*E +\n",
       "10000*S +\n",
       "0 +\n",
       "1*N +\n",
       "10*E +\n",
       "100*V +\n",
       "1000*E +\n",
       "10000*S +\n",
       "0 +\n",
       "1*E +\n",
       "10*N +\n",
       "100*I +\n",
       "1000*N ==\n",
       "0 + 1*Y + 10*T + 100*R + 1000*I + 10000*H + 100000*T"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "letters = set(expr) & set(ascii_letters)\n",
    "letter2variable = {letter: Int(letter) for letter in letters}\n",
    "\n",
    "left, right = expr.split('=')\n",
    "first, *rest = left.split() # we assume that parts are surrounded by spaces\n",
    "right = right.strip()\n",
    "\n",
    "result = word2expression(first)\n",
    "for op, word in zip(rest[::2], rest[1::2]):\n",
    "  if op == '+':\n",
    "    result += word2expression(word)\n",
    "  else:\n",
    "    result -= word2expression(word)\n",
    "\n",
    "result == word2expression(right)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Such expression, that relates variable's values to the alphametic puzzle, is just one of the constraints. Other required constraints are: \n",
    "\n",
    "* evey variable must be in the `[0, 9]` interval,\n",
    "* variables corresponding to letters that appear first in a word must be greater than `0`,\n",
    "* variables must have distinct values.\n",
    "\n",
    "We can modify the code above to compute the full list of constraints:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "letters = set(expr) & set(ascii_letters)\n",
    "letter2variable = {letter: Int(letter) for letter in letters}\n",
    "\n",
    "constr = (\n",
    "  [0 <= letter2variable[letter] for letter in letters] + \n",
    "  [letter2variable[letter] <= 9 for letter in letters]\n",
    ")\n",
    "constr.append(Distinct([letter2variable[letter] for letter in letters]))\n",
    "\n",
    "left, right = expr.split('=')\n",
    "first, *rest = left.split() # we assume that parts are surrounded by spaces\n",
    "right = right.strip()\n",
    "\n",
    "constr.extend([letter2variable[first[0]] > 0, letter2variable[right[0]] > 0])\n",
    "\n",
    "result = word2expression(first)\n",
    "for op, word in zip(rest[::2], rest[1::2]):\n",
    "  constr.append(letter2variable[word[0]] > 0)\n",
    "  if op == '+':\n",
    "    result += word2expression(word)\n",
    "  else:\n",
    "    result -= word2expression(word)\n",
    "\n",
    "constr.append(result == word2expression(right))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once all the constraints have been prepared, it is enough to ask z3 to check if they can be satisfied:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b>sat</b>"
      ],
      "text/plain": [
       "sat"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = Solver()\n",
    "s.add(constr)\n",
    "s.check()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given that the constraints are satisfiable, we can ask a model (that is an assignement of values to variables) that satisfies it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "[T = 1,\n",
       " S = 4,\n",
       " V = 7,\n",
       " H = 5,\n",
       " E = 9,\n",
       " N = 3,\n",
       " R = 6,\n",
       " I = 2,\n",
       " Y = 8]"
      ],
      "text/plain": [
       "[T = 1,\n",
       " S = 4,\n",
       " V = 7,\n",
       " H = 5,\n",
       " E = 9,\n",
       " N = 3,\n",
       " R = 6,\n",
       " I = 2,\n",
       " Y = 8]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = s.model()\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the model to write the solution in a more readable form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEVEN + SEVEN + SEVEN + NINE = THIRTY\n",
      "49793 + 49793 + 49793 + 3239 = 152618\n"
     ]
    }
   ],
   "source": [
    "print(expr + '\\n' + \n",
    "  ''.join(str(m[letter2variable[letter]]) if letter in letters else letter for letter in expr)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Putting all together in a single function we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solvealphametic(expr):\n",
    "\n",
    "  letters = set(expr) & set(ascii_letters)\n",
    "  letter2variable = {letter: Int(letter) for letter in letters}\n",
    "\n",
    "  def word2expression(word):\n",
    "    return sum(10 ** pow * letter2variable[letter] for pow, letter in enumerate(word[::-1]))\n",
    "\n",
    "  constr = (\n",
    "    [0 <= letter2variable[letter] for letter in letters] + \n",
    "    [letter2variable[letter] <= 9 for letter in letters]\n",
    "  )\n",
    "  constr.append(Distinct([letter2variable[letter] for letter in letters]))\n",
    "\n",
    "  left, right = expr.split('=')\n",
    "  first, *rest = left.split() # we assume that parts are surrounded by spaces\n",
    "  right = right.strip()\n",
    "\n",
    "  constr.extend([letter2variable[first[0]] > 0, letter2variable[right[0]] > 0])\n",
    "\n",
    "  result = word2expression(first)\n",
    "  for op, word in zip(rest[::2], rest[1::2]):\n",
    "    constr.append(letter2variable[word[0]] > 0)\n",
    "    if op == '+':\n",
    "      result += word2expression(word)\n",
    "    else:\n",
    "      result -= word2expression(word)\n",
    "\n",
    "  constr.append(result == word2expression(right))\n",
    "\n",
    "  s = Solver()\n",
    "  s.add(constr)\n",
    "  if s.check() == sat:\n",
    "    m = s.model()\n",
    "    print(expr + '\\n' + \n",
    "      ''.join(str(m[letter2variable[letter]]) if letter in letters else letter for letter in expr)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's time to go back to the first example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEND + MORE = MONEY\n",
      "9567 + 1085 = 10652\n"
     ]
    }
   ],
   "source": [
    "solvealphametic('SEND + MORE = MONEY')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few other examples are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "terrible + number = thirteen\n",
      "45881795 + 302758 = 46184553\n"
     ]
    }
   ],
   "source": [
    "solvealphametic('terrible + number = thirteen')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "romans + also + more + or + less + added = letters\n",
      "975348 + 3187 + 5790 + 79 + 1088 + 36606 = 1022098\n"
     ]
    }
   ],
   "source": [
    "solvealphametic('romans + also + more + or + less + added = letters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NINETEEN + THIRTEEN + THREE + TWO + TWO + ONE + ONE + ONE = FORTYTWO\n",
      "42415114 + 56275114 + 56711 + 538 + 538 + 841 + 841 + 841 = 98750538\n"
     ]
    }
   ],
   "source": [
    "solvealphametic('NINETEEN + THIRTEEN + THREE + TWO + TWO + ONE + ONE + ONE = FORTYTWO')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "86f1961004bc775655e233a19f843927dfe29b0169db43be7fbad477fb6a4a49"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit ('b812212f80584f228cf2fbd10554d7ac-jupyter-fun': venv)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}