Wednesday, July 13, 2011

DFS for puzzle solving

I looked back at the one to nine puzzle, and realized that we can write a function to hold the main idea, and apply it to other puzzles.

The main idea is to solve the puzzle one step at a time. At each step, we try every possible move, but only the moves that are still valid. We return any solutions we find.

In the previous post, the searching was mixed with the puzzle specific code. Here is a function that handles the recursion for us, and lets us pass in problem-specific code for the details.

let solve next_f done_f initial =
    let rec search state =
        seq {
            if done_f state then
               yield state
            else
               for state' in next_f state do
                   yield! search state'
            }
    search initial

This function returns a sequence of answers. Some puzzles have one solution, others have many. With sequences, the function caller has more control over how many values are returned.

done_f has the type ('position -> bool) and is the stopping criteria. Our move generator next_f has the type ('position -> seq<'position>). solve is a higher order function because it takes other functions as inputs.

Now, we can solve the one to nine puzzle with this function.

let one_to_nine () =
    let next_move (number, digits, idx) =
        let idx' = idx + 1
        seq {
            for d in digits do
               let number' = number * 10 + d
               if number' % idx' = 0 then
                   yield number', Set.remove d digits, idx'
            }

    let finished (_, _, idx) = idx = 9

    solve next_move finished (0, Set.ofList [1..9], 0)
    |> Seq.head 
    |> fun (number, digits, idx) -> number


see the code at fssnip.net

A position is three things: the current number, the available digits, and the index of the digit we are considering. If we make it to nine digits, then we are done. Given a current state, we generate the next state by trying each digit that remains. We test each of these to make sure they follow the divisibility rule (two digit number must be divisible by two, three by three, etc.) Then we create a new state to hold the information.

We pass those two functions to solve along with the initial state, and it finds our solution.

solve is meant to be general, and that was too much work for one puzzle. Let's look next at the N-Queens problem. There are a ton of possible positions we could try for all N queens, but each piece reduces the number of options we have for the other pieces. We will place the queens one column at a time. At each column, we will compute which rows are valid based on the previous columns. A row is valid if no previous column has a queen on the same row or diagonal.

N-Queens has many solutions. With solve, we can find one, some, or all of the solutions, given enough time or memory.

Here is the code.

let n_queens n =
    let all_rows = Set.ofList [1..n]
    let next_f queens =
        let column = List.length queens + 1
        queens
        |> List.fold (fun possible (row, col) ->
            let distance = column - col
            // Remove the current row, and diagonals from the
            // possible set.  The diagonals are just previous
            // row +/- the column difference. This is safe here,
            // because even if row - distance is negative, we are
            // removing it from another set.  No array access is
            // involved.
            Set.difference possible
                      (Set.ofList [row; row + distance;
                                   row - distance])) all_rows
        |> Set.map (fun (row) ->
                (row, column)::queens)
    let done_f queens =
        List.length queens = n

    solve next_f done_f []

see the code at fssnip.net

And some results from running the code.

n_queens 8 |> Seq.length;;
val it : int = 92

> n_queens 9 |> Seq.length;;
val it : int = 352

> n_queens 10 |> Seq.length;;
val it : int = 724

> n_queens 11 |> Seq.length;;
val it : int = 2680

> // Solve for N = 20, but just return one solution. 
> n_queens 20 |> Seq.head;;
val it : (int * int) list =
  [(11, 20); (6, 19); (14, 18); (7, 17); (10, 16); (8, 15); (19, 14); (16, 13);
   (9, 12); (17, 11); (20, 10); (18, 9); (12, 8); (15, 7); (13, 6); (4, 5);
   (2, 4); (5, 3); (3, 2); (1, 1)]

This function also solves Sudoku, with the right inputs. By the time all of the helper code is written, solve is a small fraction of the code. But, it is still instructive to prune away the details of the puzzles to find the essence of the search problem, even if only for better understanding.

Sudoku code