Latin square solver in clojure

latinsq.clj

(ns latinsq.core
  (:use [clojure.set :only (difference)]))


(defn replace-at
  "in string s, replaces character at index p with c"
  [s p c]
  (str
   (.substring s 0 p)
   c
   (.substring s (inc p))))

; memoized function to save time on sqrt calls
(def sqrt (memoize (fn [x] (Math/sqrt x))))


(defn candidates
  "assuming that the position pos is empty in sq, return those values that it might be"
  [pos sq]
  
  (let [sq-size (int (sqrt (count sq)))]

					; set difference between...
    (difference
					; ...all the possible values...
     (set (map #(first (str %)) (range 1 (inc sq-size))))

					; ...and the set of...
     (into #{}
	   (concat
					; ...those in the same column...
	    (map #(get sq %)
		 (range (rem pos sq-size)
			(count sq)
			sq-size))
					; ...and those in the same row.
	    (map #(get sq %)
		 (map #(+ % (- pos (rem pos sq-size)))
		      (range 0 sq-size))))))))




(defn latinsq
  "encode your partial-square as a string like 1--1
   this fn returns a lazy sequence of all solutions"
  [sq]
					; Find the first empty square
  (let [empty-pos (.indexOf sq "-")]

					; if none, we don't need to do anything
    (if (= -1 empty-pos)
      (list sq)

					; else make a lazy sequence of...
      (lazy-seq
					; ...the concatenation of all the results of...
       (apply concat
					; ...using "map" to recurse, filling in the empty
					; square with...
	      (map #(latinsq (replace-at sq empty-pos %))

					; ...each possible value in turn
		   (candidates empty-pos sq)))))))




;; So, now some examples

(time
 (latinsq "123------"))
;; "Elapsed time: 0.045368 msecs"
;; ("123231312" "123312231")

(time
 (latinsq "12---31--------4"))
;; "Elapsed time: 0.068511 msecs"
;; ("1243431224313124" "1243431234212134")


;; A bit harder, an empty 5x5 grid
;; has 161280 solutions according to
;; http://mathworld.wolfram.com/LatinSquare.html
(time
 (count (latinsq "-------------------------")))
;; "Elapsed time: 36980.759177 msecs"   <--- a bit slow
;; 161280


;; Having made sure that our function returns a lazy seq
;; the whole result can be treated lazily, so to find just one
;; solution to the 5x5 grid:
(time
 (first (latinsq "-------------------------")))
;; "Elapsed time: 0.985559 msecs"       <--- not slow
;; "1234521453345124523153124"

;; first 3 of 5524751496156892842531225600 solutions for a 9x9 grid
(time 
 (take 3 (latinsq "---------------------------------------------------------------------------------")))
;; "Elapsed time: 0.075874 msecs"
;; ("123456789214365897341278956432189675567891234658917342789523461896742513975634128"
;;  "123456789214365897341278956432189675567891234658917342789523461975634128896742513"
;;  "123456789214365897341278956432189675567891234658917342789524163896743521975632418")