Killer Sudoku solver

sudoku.kt

import java.io.BufferedWriter
import java.io.OutputStreamWriter

typealias Grid = List<List<Int>>

data class Cell(val row: Int, val col: Int)

data class Group(val label: String, val total: Int? = null, val cells: List<Cell>)

data class NewBoard(val n: Int,
                    val cells: List<Cell>,
                    val values: Map<Cell, Int>,
                    val groups: List<Group>) {
    companion object {
        fun sudoku(grid: Grid): NewBoard {
            val n = grid.size
            val cellValues = grid
                    .flatMap { x -> x.asIterable() }
                    .mapIndexed { index, value ->
                        val row = index / n
                        val col = index % n
                        Pair(Cell(row, col), value)
                    }
            val cells = cellValues.map { it -> it.first }
            val values = cellValues.asSequence().filter { it.second != 0 }
                    .associateBy({ it -> it.first }, { it -> it.second })

            val rowGroups = cells.asSequence()
                    .groupBy { it -> it.row }
                    .map { (row, cells) -> Group("Row $row", null, cells) }
                    .toList()

            val colGroups = cells.asSequence()
                    .groupBy { it -> it.col }
                    .map { (col, cells) -> Group("Col $col", null, cells) }
                    .toList()

            val boxGroups = if (n == 9) {
                cells.asSequence()
                        .groupBy { cell -> Pair(cell.row / 3, cell.col / 3) }
                        .map { it -> Group("Box (${it.key.first},${it.key.second})", null, it.value) }
                        .toList()
            } else emptyList()

            return NewBoard(n, cells, values, rowGroups + colGroups + boxGroups)
        }


        /**
         * Killer Sudoku, the board is empty, but there are extra constraint groups with totals
         * that must be satisfied
         */
        fun killer(n: Int, constraints: List<Group>) : NewBoard {
            val empty = List(n) { _ -> List(n) {0} }
            val board = sudoku(empty)
            val cCells : Set<Cell> = constraints.flatMap { g -> g.cells.asIterable() }.toSet()
            assert(cCells.size == board.cells.size)
            return NewBoard(board.n, board.cells, board.values, board.groups + constraints)
        }
    }

    fun print(writer: java.io.BufferedWriter) {
        var row = 0
        writer.newLine()
        cells.forEach {
            if (row != it.row) {
                writer.newLine()
                row++
            }
            writer.write(" ${values.getOrDefault(it, 0)}")
        }
        writer.newLine()
        writer.flush()
    }

    fun guessCell(cell: Cell, value: Int): NewBoard {
        val newValues = values.toMutableMap().apply { put(cell, value) }
        return NewBoard(n, cells, newValues, groups)
    }

    fun isValid(): Boolean {
        return groups.fold(true) { valid, g -> valid && groupValid(g) }
    }

    private fun groupValid(group: Group): Boolean {
        val entries = group.cells.asSequence().map { it -> values[it] }.filterNotNull().toList()

        return entries.size == entries.toSet().size &&
                (group.total == null || if (entries.size < group.cells.size) {
                                            (entries.sum() < group.total)
                                        } else {
                                            (entries.sum() == group.total)
                                        })

    }
}

fun solve(board: NewBoard, index: Int = 0): NewBoard? {
    //board.print(BufferedWriter(OutputStreamWriter(System.out)))

    if (!board.isValid()) {
        return null
    } else if(index == board.cells.size) {
        return board
    } else {
        val cell = board.cells[index]

        if (board.values.containsKey(cell)) {
            return solve(board, index + 1)
        } else {
            for (i in 1..board.n) {
                solve(board.guessCell(cell, i), index + 1)?.let { return it }
            }
        }
    }
    return null
}

fun main(args: Array<String>) {
    val writer = BufferedWriter(OutputStreamWriter(System.out))

    val board: NewBoard = NewBoard.sudoku(listOf(
            listOf(4, 0, 1, 0, 0, 3),
            listOf(0, 0, 0, 1, 0, 0),
            listOf(0, 4, 2, 3, 0, 6),
            listOf(1, 0, 3, 2, 5, 0),
            listOf(0, 0, 5, 0, 0, 0),
            listOf(6, 0, 0, 5, 0, 1)))
    writer.write("\nSolving simple sudoku\n")
    solve(board)?.print(writer)

    val board2: NewBoard = NewBoard.sudoku(listOf(
            listOf(0, 0, 0, 0, 0, 0, 2, 0, 0),
            listOf(0, 0, 0, 0, 2, 0, 0, 3, 0),
            listOf(2, 0, 5, 1, 9, 0, 0, 0, 6),
            listOf(0, 0, 0, 0, 1, 4, 0, 6, 3),
            listOf(0, 3, 0, 0, 0, 0, 0, 7, 0),
            listOf(4, 2, 0, 9, 3 ,0, 0, 0, 0),
            listOf(1, 0, 0, 0, 6, 2, 5, 0, 9),
            listOf(0, 6, 0, 0, 5, 0, 0, 0, 0),
            listOf(0, 0, 4, 0, 0 ,0, 0, 0, 0)))

    writer.write("\nSolving normal sudoku\n")

    solve(board2)?.print(writer)


    // TODO: better import format to represent the groups..
    val board3: NewBoard = NewBoard.killer(9, listOf(
            Group("0,0", 14, listOf(Cell(0, 0),
                                    Cell(1, 0))),
            Group("0,1", 13, listOf(Cell(0, 1), Cell(0, 2))),
            Group("0,3",  9, listOf(Cell(0, 3), Cell(0, 4))),
            Group("0,5", 14, listOf(Cell(0, 5), Cell(0, 6),
                                    Cell(1, 5))),
            Group("0,7",  8, listOf(Cell(0, 7), Cell(0, 8))),

            //ROW2
            Group("1,1",  5, listOf(Cell(1, 1), Cell(1, 2))),
            Group("1,3", 23, listOf(Cell(1, 3), Cell(1, 4),
                                                Cell(2, 4))),
            Group("1,6", 22, listOf(Cell(1, 6), Cell(1, 7),
                                    Cell(2, 6), Cell(2, 7))),
            Group("1,8",  7, listOf(Cell(1, 8), Cell(2, 8))),


            Group("2,0",  4, listOf(Cell(2, 0),
                                    Cell(3, 0))),
            Group("2,1", 12, listOf(Cell(2, 1), Cell(2, 2))),
            Group("2,3",  9, listOf(Cell(2, 3), Cell(3, 3))),
            Group("2,5",  8, listOf(            Cell(2, 5),
                                    Cell(3, 4), Cell(3, 5))),


            Group("3,1", 18, listOf(Cell(3, 1), Cell(3, 2),
                                    Cell(4, 1), Cell(4, 2))),
            Group("3,6", 23, listOf(Cell(3, 6), Cell(3, 7), Cell(3, 8))),


            Group("4,0", 20, listOf(Cell(4, 0), Cell(5, 0), Cell(6, 0))),
            Group("4,3", 17, listOf(Cell(4, 3), Cell(5, 3), Cell(6, 3))),
            Group("4,4",  5, listOf(Cell(4, 4), Cell(5, 4))),
            Group("4,5", 21, listOf(Cell(4, 5), Cell(4, 6),
                                    Cell(5, 5))),
            Group("4,7",  7, listOf(Cell(4, 7), Cell(4, 8))),


            Group("5,1", 11, listOf(Cell(5, 1), Cell(5, 2))),
            Group("5,6",  3, listOf(Cell(5, 6), Cell(5, 7))),
            Group("5,8", 16, listOf(Cell(5, 8),
                                    Cell(6, 8))),


            Group("6,1", 17, listOf(Cell(6, 1),
                                    Cell(7, 1))),
            Group("6,2",  4, listOf(Cell(6, 2),
                                    Cell(7, 2))),
            Group("6,4", 12, listOf(            Cell(6, 4), Cell(6, 5),
                                    Cell(7, 3), Cell(7, 4))),
            Group("6,6", 11, listOf(Cell(6, 6), Cell(6, 7))),


            Group("7,0",  7, listOf(Cell(7, 0),
                                    Cell(8, 0))),
            Group("7,5", 14, listOf(Cell(7, 5),
                                    Cell(8, 5))),
            Group("7,6", 15, listOf(Cell(7, 6), Cell(7, 7))),
            Group("7,8",  6, listOf(Cell(7, 8),
                                    Cell(8, 8))),

            Group("8,1", 10, listOf(Cell(8, 1), Cell(8, 2))),
            Group("8,3", 16, listOf(Cell(8, 3), Cell(8, 4))),
            Group("8,6",  4, listOf(Cell(8, 6), Cell(8, 7)))

    ))

    writer.write("\nSolving killer sudoku\n")
    solve(board3)?.print(writer)

}