Thursday 3 November 2011

A chess problem solver

Some time ago when I was looking for a job, basically every company I applied to had some kind of preliminary programming exercise that one needs to pass before getting to a live interview. There was quite a bit of variance in these exercises. The most interesting (and the most challenging) by far was a chess problem solver.

In this problem, you're given a chess board of size n by m and a bunch of pieces excluding pawns. It is assumed that all the pieces are colourless anarchists that ceaselessly threaten each other when given the chance. The task: print out all the ways the pieces can be placed on the board so that they don't threaten each other.

I decided to do this one in Clojure, which is currently my favourite language. So let's see how that went.

(Notice that the following code might contain severe performance-related suboptimalities. While the code isn't very optimised, it worked well enough to get me the interview. More on the performance later.)

Before we begin, let's declare the namespace and include some stuff we'll need later:

(ns chess.core
  (:require [clojure.string :as str]
            [clojure.set :as set]
            [clojure.contrib.combinatorics :as combi]))

The first thing to decide is the representation. For some reason, I didn't represent the board explicitly; instead I went with a hash map of [x y] -> piece for the piece locations and a separate vector [bx by] for the board size. A piece is just its name - a keyword: :Q for queen, :K for king, :B for bishop, :N for knight, :R for rook.

The first thing to code then was something that takes a board - that is, a board-size vector and a location-to-piece hash map - and prints it in a more human-readable format:

(defn get-board [[bx by] loc-pieces]
  (partition bx
             (for [j (range by)
                   i (range bx)] ;; x first
               (when-let [a (get loc-pieces [i j])]
                 (name a))))) ;; when :X => "X", otherwise nil

(defn print-board [board]
  (str/join "\n" (for [row board]
                   (str/join "" (for [cell row]
                                  (if cell cell "."))))))

(defn pb [bs locs]
  (-> (get-board bs locs) (print-board) (println)))

For instance:

chess.core> (pb [3 5] {[0 0] :R, [1 1] :N, [2 2] :R, [1 3] :N, [1 4] :B})
R..
.N.
..R
.N.
.B.
nil
chess.core>

Now the algorithm itself is just a variation of basic depth-first search. The idea is simple: for each piece, we maintain a set of the free squares on which that piece can be placed. When placing a new piece, this set is updated by removing those squares that are now impossible. If the set becomes empty for any piece, the situation is impossible, so we backtrack.

To make this more precise, here's the algorithm in pseudocode:

def solve(boardSize, locPieceMap, freeSquares, remPieces):
    if remPieces is empty:
        return locPieceMap
    moves = getMovesForBestPiece(boardSize, locPieceMap, freeSquares, remPieces)
    if moves is empty:
        return None
    solutions = []
    for piece, locMapNow, freesNow in moves:
        recResult = solve(boardSize, locMapNow, freesNow, remPieces-piece)
        if recResult:
            solutions.append(recResult)
    return solutions

Here getMovesForBestPiece is a function that figures out which piece we should play next, using some heuristic. The function returns a list of valid moves for the best piece. It also checks for impossible situations such as some piece not having any suitable squares; in this case the function returns empty.

What's a good heuristic? I tried three different ones:
  1. Choose the piece of which we have the fewest. E.g. if we have 1 queen and 2 knights left to place, choose the queen.
  2. Choose the piece whose free-squares count is the lowest. E.g. if we have 1 bishop which has 11 squares it can potentially be placed on, and 2 kings which each have 8 potential squares, choose the kings.
  3. Try each piece on each square it can go on, and choose the piece which results in the fewest options after that.
Not surprisingly, option 3 was hilariously slow and thus abandoned. The first two are much faster. In my test runs they seem to be within 2% of each other; I'm not sure why this is, but hey.

Before looking at the details of getMoves, let's see how to deal with duplicate avoidance. When you have multiple copies of the same piece, a naive depth-first search implementation will give you duplicates. Consider the situation where there are, say, two free squares A and B and two pieces K1 and K2. The pieces are both kings and thus identical, but the algorithm will consider the cases "K1 on A; K2 on B" and "K2 on A; K1 on B" as separate, unless instructed otherwise.

To avoid duplicates, I decided to place all identical pieces on the board at the same time. (Another option would have been to maintain a list of already used squares for the current piece type, which is essentially the same thing, except more complicated.) I simply create each combination of (free-squares choose number-of-pieces-to-place) and iterate through those. The combinations are generated with the combinations function from clojure.contrib.combinatorics; this is lazy and so probably efficient enough - it is necessary to iterate through every combination, but we don't want to pre-calculate too many.

So let's look at what getMoves needs to do, in pseudocode.

def getMovesForBestPiece(boardSize, locPieceMap, freeSquares, remPieces):
    piece = guessBestPiece(boardSize, locPieceMap, freeSquares, remPieces)
    frees = freeSquares[piece]
    count = remPieces[piece]
    moves = []
    for squares in combinations(frees, count):
        locPieceMapNow = locPieceMap+{squares:piece}
        newThreateneds = squaresThreatenedBy(locPieceMapNow)
        # cannot place pieces so that they threaten any other piece
        if intersection(newThreateneds, locPieceMapNow) is not empty:
            continue
        nonFreeNow = union(locPieceMapNow, newThreateneds)
        freeNow = updateFreeSquares(boardSize, locPieceMapNow, freeSquares, nonFreeNow, remPieces)
        # require valid squares for each piece
        if any in freeNow is empty:
            continue
        moves.append([piece, locPieceMapNow, freeNow])
    return moves

(Note that in case of heuristic 3 above, the code would be slightly different: instead of guessing the piece, we'd run this routine for every remaining piece and then choose.)

Fairly straightforward, I guess. The set intersection and union are found in clojure.set, so all we need now are squaresThreatenedBy and updateFreeSquares. For these I'll go straight to the implementation.

(defn remove-nonexistent-squares [squares [bx by]]
  (for [[x y] squares :when (and (>= x 0) (>= y 0) (< x bx) (< y by))]
    [x y]))

(defn squares-threatened-by [[location piece] board-size]
  (let [[x y] location
        [bx by] board-size]
    (condp = piece
      :K (let [neighbours (remove (partial = [0 0]) (for [i [-1 0 1] j [-1 0 1]] [i j]))
               king-squares (for [[nx ny] neighbours]
                              [(+ x nx) (+ y ny)])]
           (remove-nonexistent-squares king-squares board-size))
      :R (let [lines (into (for [i (range bx)] [i y]) (for [j (range by)] [x j]))
               rook-squares (remove (partial = location) lines)]
           rook-squares) ; construction doesn't generate nonexistents
      :N (let [xs [1 1 2 2 -1 -1 -2 -2]
               ys [2 -2 1 -1 2 -2 1 -1]
               coord-adds (map list xs ys)
               knight-squares (for [[xa ya] coord-adds]
                                [(+ x xa) (+ y ya)])]
           (remove-nonexistent-squares knight-squares board-size))
      :B (let [coord-adds (for [i (range 1 (max bx by))] [i i])
               bishop-squares (apply concat
                                     (for [[xa ya] coord-adds]
                                       (map #(list (%1 x xa) (%2 y ya)) [+ + - -] [+ - + -])))]
           (remove-nonexistent-squares bishop-squares board-size))
      :Q (let [rooks (squares-threatened-by [location :R] board-size)
               bishs (squares-threatened-by [location :B] board-size)]
           (into rooks bishs)) ;; rooks, bishs are distinct by construction
      ;; else:
      (str "Error: unknown piece " piece))))

That was fun. Let's try it out:

chess.core> (squares-threatened-by [[2 2] :Q] [3 4])
([0 0] [1 1] [1 3] [2 3] [2 1] [2 0] [0 2] [1 2])
chess.core> (squares-threatened-by [[0 0] :K] [3 3])
([0 1] [1 0] [1 1])
chess.core> (squares-threatened-by [[2 1] :R] [3 3])
([2 2] [2 0] [0 1] [1 1])
chess.core> (squares-threatened-by [[1 2] :N] [4 4])
([2 0] [3 3] [3 1] [0 0])
chess.core> (squares-threatened-by [[1 2] :B] [4 5])
([2 3] [2 1] [0 3] [0 1] [3 4] [3 0])
chess.core>

Seems to work, as you can verify.

For the free squares, the data structure is a map where each piece is associated with a set of the squares that are yet free for it. To illustrate, the initial setup is one where each piece is associated with all the squares:

(defn get-initial-free-squares [[bx by] piece-types]
  (let [all-squares (set (for [i (range bx), j (range by)] [i j]))]
    (into {} (for [piece piece-types] {piece all-squares}))))

chess.core> (get-initial-free-squares [2 2] [:Q :R])
{:Q #{[1 0] [0 0] [1 1] [0 1]}, :R #{[1 0] [0 0] [1 1] [0 1]}}
chess.core>

With this structure, the code for update-free-squares is then:

(defn update-free-squares [board-size loc-pieces free-squares non-frees piece-types]
  (into {}
        (for [[piece frees] free-squares
              :let [rem-frees (set/difference frees non-frees)]]
          {piece rem-frees})))

Now, armed with squares-threatened-by and update-free-squares, we can finally implement the real version of get-moves (here with heuristic #1):

(defn get-piece-square-combinations [piece squares piece-count]
  (let [pieces (repeat piece-count piece)
        square-combs (combi/combinations squares piece-count)]
    (for [comb square-combs]
      (zipmap comb pieces))))

(defn get-moves-for-piece [board-size loc-pieces free-squares piece-counts piece]
  (let [frees (get free-squares piece)
        piece-count (get piece-counts piece)
        piece-square-combos (get-piece-square-combinations piece frees piece-count)
        rem-piece-counts (dissoc piece-counts piece)
        rem-free-squares (dissoc free-squares piece)]
    (for [attempt piece-square-combos
          :let [squares (keys attempt)
                locs-now (into loc-pieces attempt)
                loc-set (set (keys locs-now))
                threatened-now (set (apply concat (for [loc-piece attempt]
                                                    (squares-threatened-by
                                                     loc-piece board-size))))]
          ;; cannot place pieces so they threaten each other/existing pieces
          :when (empty? (set/intersection loc-set threatened-now))
          :let [non-free-now (set/union loc-set threatened-now)
                free-now (update-free-squares board-size locs-now rem-free-squares
                                              non-free-now (keys rem-piece-counts))]
          ;; require valid squares for all remaining pieces
          :when (every? pos? (map count (vals free-now)))]
      [piece locs-now free-now])))

(defn get-moves [board-size loc-pieces free-squares piece-counts]
  (if (some empty? (vals free-squares))
    ()
    (let [piece (first (apply min-key second piece-counts))]
      (get-moves-for-piece board-size loc-pieces free-squares piece-counts piece))))

All that's left now is solve. But before looking at that, there's an optimisation which turns out to be beneficial.

In algorithms like this, there's a balance of how much to compute on each step. If you do too much work on each round, as in heuristic #3, the rounds take too long and the algorithm is slow; but if you do too little, you end up recursing into impossible paths too much, which is also slow. So to detect impossible paths earlier, it can be a good idea to lift some of the calculations for the next round into this round, as long as you don't overdo it and end up with too-heavy rounds.

In our case, it turns out we can optimise get-moves a bit to make it detect impossible cases earlier. For example, suppose we're trying to fit a rook, a queen and a knight on a 3x3 board. The situation might be that we have placed the rook on [2, 1], leaving us with 4 free squares for each remaining piece, and our heuristic randomly recommends that we place the knight next. We then run get-moves-for-piece for the knight. Here's the situation:

chess.core> (pb [3 3] {[2 1] :R})
...
..R
...
nil
chess.core> (def frees {:Q #{[1 0] [0 0] [1 2] [0 2]}, :N #{[1 0] [0 0] [1 2] [0 2]}})
#'chess.core/frees
chess.core> (get-moves-for-piece [3 3] {[2 1] :R} frees {:Q 1, :N 1} :N)
;;; a list of: [piece placed; board after placing; free squares after placing]
([:N {[2 1] :R, [1 0] :N} {:Q #{[0 0] [1 2]}}]
 [:N {[2 1] :R, [1 2] :N} {:Q #{[1 0] [0 2]}}])
chess.core>

Looking at the boards resulting from these moves, you can see that in fact the queen cannot be placed on either of its potentially free squares:

chess.core> (pb [3 3] {[2 1] :R, [1 0] :N})
.N.
..R
...
nil
chess.core> (pb [3 3] {[2 1] :R, [1 2] :N})
...
..R
.N.
nil
chess.core>

Now these free squares are just fine for, say, a bishop, but they're not compatible with a queen. On the next round, the algorithm will notice this incompatibility and backtrack. However, we can lift the computation that notices this into this round, and this happens to be beneficial in this case, resulting in a much faster algorithm.

Here's the uplifted version of the code. In update-free-squares, just the last line is changed to use the new pre-calculation function.

(defn remove-threatening-squares [board-size loc-pieces free-square-set piece]
  (set
   (let [locs (set (keys loc-pieces))]
     (for [square free-square-set
           :let [threateneds (set (squares-threatened-by [square piece] board-size))]
           :when (empty? (set/intersection locs threateneds))]
       square))))

(defn update-free-squares [board-size loc-pieces free-squares non-frees piece-types]
  (into {}
        (for [[piece frees] free-squares
              :let [rem-frees (set/difference frees non-frees)]]
          {piece (remove-threatening-squares board-size loc-pieces rem-frees piece)})))

The remove-threatening-squares function goes through each preliminarily free square for the given piece. When that piece placed on the square results in one of the existing pieces being threatened, that square is filtered out. Like so:

chess.core> (get-moves-for-piece [3 3] {[2 1] :R} frees {:Q 1, :N 1} :N)
()
chess.core>

That's better.

To wrap up, here's the solve function, plus related auxiliary stuff.

(defn handle-solution [board-size loc-pieces]
  (pb board-size loc-pieces)
  (println)
  [1]) ;; could also return the solution

(defn solve-rec [board-size loc-pieces free-squares piece-counts]
  (if (empty? piece-counts)
    (handle-solution board-size loc-pieces)
    (let [moves (get-moves board-size loc-pieces free-squares piece-counts)]
      (if (empty? moves)
        ()
        (apply concat
               (for [[piece locs-now frees-now] moves
                     :let [rem-counts (dissoc piece-counts piece)
                           rec-result (solve-rec board-size locs-now frees-now rem-counts)]
                     :when (not (empty? rec-result))]
                 rec-result))))))

(defn solve [board-size piece-counts]
  (let [initial-frees (get-initial-free-squares board-size (keys piece-counts))
        result (solve-rec board-size {} initial-frees piece-counts)
        solution-count (count result)]
    (println (str "\n" "total number of solutions: " solution-count "\n"))
    solution-count))

Phew. A test run:

chess.core> (solve [3 3] {:K 1, :R 2})
.R.
R..
..K

K..
..R
.R.

.R.
..R
K..

..K
R..
.R.


total number of solutions: 4

4
chess.core>

And that seems to be that!

Update: the code is now available on GitHub. The version there includes some extra code for debugging and all 3 heuristics. Also available are the unit tests, and a Leiningen project to run all that stuff a bit more easily.

Finally, a note on performance. The specification for this exercise gave a rather large example and asked for the exact number of solutions for it; in addition, the code had to be capable of enumerating these solutions within one hour. The code just presented in fact took longer than that (but not much longer). I haven't profiled my code or, indeed, given its performance much thought, since it turned out to be close enough. Proper optimisation of this is a task for another day.

If you want to test the performance yourself, here's an example of similar size; the goal is to get this to run in less than an hour:

;;; from a REPL:
(solve [7 8] {:K 3, :Q 1, :B 2, :R 2, :N 3})

# or from the command line, with the script included in the package:
./run '[7 8]' '{:K 3, :Q 1, :B 2, :R 2, :N 3}'

How fast would this algorithm run when properly optimised? Or when implemented in, say, Python or Java or C? Or even better, what would be a more efficient algorithm?

1 comment: