ぷよぷよ19連鎖の問題を Clojure で

去年、迷路を最短経路で解く問題を出題していたところが、今年も採用一次試験の問題を公開しているというのでClojureで解いてみた。

問題は以下のとおり。

ゲーム「ぷよぷよ」で、フィールドの状態がテキストで与えられたとき、消える「ぷよ」を消して次のフィールドの状態を出力するプログラムを書け。
たとえば、色をG/Y/Rで表すとき(Green/Yellow/Red)、

GGR
YGG

であればGが消えて

Y R

になります。

また、このプログラムを使って次のフィールドを与えると19連鎖ののちすべてのぷよが消えることを確認し、消える途中の様子をあわせて提出すること。

  GYRR
RYYGYG
GYGYRR
RYGYRG
YGYRYG
GYRYRG
YGYRYR
YGYRYR
YRRGRG
RYGYGG
GRYGYR
GRYGYR
GRYGYR 

人生を書き換える者すらいた。: 人材募集企画 2011年版

ソースコード

puyo.clj

(use '[clojure.contrib.io :only (read-lines)]
     '[clojure.string :only (join)])

(defn value-at [field y x]
  (get-in field [x y]))

(defn nrows [field]
  (count (field 0)))

(defn ncols [field]
  (count field))

(defn entire-block
  ([pos field]
   (entire-block pos field #{}))
  ([[y x :as pos] field visited]
   (reduce (fn [visited* pos*]
             (if (visited* pos*)
               visited*
               (into visited* (entire-block pos* field visited*))))
           (conj visited pos)
           (for [[dy dx] [[0 -1] [1 0] [0 1]]
                 :let [[y* x* :as pos*] [(+ y dy) (+ x dx)]]
                 :when (and (not (visited pos*))
                            (= (value-at field y x)
                               (value-at field y* x*)))]
             pos*))))

(defn all-blocks [field]
  (let [blocks (atom [])]
    (doseq [y (range (nrows field))
            x (range (ncols field))
            :let [pos [y x]]
            :when (and (not= (value-at field y x) \space)
                       (not (some #(% pos) @blocks)))
            :let [block (entire-block pos field)]
            :when (>= (count block) 4)]
      (swap! blocks conj block))
    @blocks))

(defn remove-blocks [field blocks]
  (let [to-be-removed (reduce into blocks)]
    (vec (map (fn [col] (vec (sort (fn [x y] (= x \space)) col)))
              (for [x (range (ncols field))]
                (for [y (range (nrows field))]
                  (if (to-be-removed [y x])
                    \space
                    (value-at field y x))))))))

(defn steps [field]
  (lazy-seq
    (cons field
          (let [blocks (all-blocks field)]
            (if-not (empty? blocks)
              (steps (remove-blocks field blocks)))))))

(defn print-field [field]
  (let [trans (apply map list field)
        color #(str "\033[" (or ({\R 41, \G 44, \Y "1;43"} %) 0) "m" %)
        hline (join (repeat (+ (ncols field) 2) "-"))
        lines `(~hline
                ~@(map #(join `("|" ~@(map color %) "\033[0m|")) trans)
                ~hline)]
    (println (join \newline lines))))

(defn solve [field]
  (doseq [step (steps field)]
    (print "\033[2J")
    (print-field step)
    (Thread/sleep 1000)))

(solve (vec (apply map vector (read-lines *in*))))

実行

標準入力からぷよの配置を問題のフォーマットで読み込み、1秒毎に各時点でのフィールドの状態を表示します。表示は見やすいように色を付けてます。緑はそのままだと黄色との境界が分かりづらかったので青で表示しています。

実行結果は下のような感じになります。