【ABC469C】「当たりが出たらもう一回」とは?

関連記事

1. 問題

`o`/`x` からなる長さ NN の文字列 SS が与えられる。
k=1,,Nk=1,\dots,N それぞれについて次を求める。
先頭 kk 個の袋を受け取って中身を食べ、以後「持っている `o` の袋を 1 個捨てて、列の先頭の袋を 1 個受け取って食べる」を、列が空になるか `o` の袋が尽きるまで繰り返す。
食べた個数(=受け取った袋の総数)を NN 行で出力せよ。

制約は 1N8×1051\le N\le 8\times10^5

2. 素朴なコード(4 TLE, MLE)

まずは、素朴に解いてみました。

(defun number-eat (str k)
  (cond ((= k 0) 0)
        ((> k (length str)) (length str))
        (t (+ k (number-eat (right str k) (contains #\o (left str k)))))))

(defun right (str n)
  (subseq str n))

(defun left (str n)
  (subseq str 0 n))

(defun contains (ch str)
  (count-if (lambda (c) (char= c ch)) str))


(defun main ()
  (let ((n (read))
        (str (read-line)))
    (loop for k from 1 to n
          do (princ (number-eat str k))
             (terpri))))

(main)
Code language: Lisp (lisp)

AC 4、TLE 12、MLE 2。

2. 素朴なコード(4 TLE, MLE)

Nの最大値は 10810^8
大きな数字になると再計算が多く、時間がかかってしまいます。
また、メモリは再帰関数の深さにも問題がありそうです。

3. 逆に考えると(AC18 418ms)

「はじめに k 個もらって、その中に “o”があれば、さらに追加してもらえる」というのは、当たりが連鎖的に出てくる状況を考える必要があります。

しかし、”x”を基準に考えると、別の見方ができます。
「はじめに k 個のストックがあり、”x”を引くごとに1ずつ減り、0 になるまで進める」。

このように考えると、先頭からの”x”の個数を数えて、それが k 個目になる位置が答えになります。
そうすると、”x”の位置を記録すれば、それが答えに相当します。

;; oxoxo ; 2 4 5
(defun miss-positions (str)
  (loop for idx below (length str)
        for ch across str
        when (char= ch #\x)
          collect (1+ idx)))

(defun main/2 ()
  (let ((n (read))
        (lst (miss-positions (read-line))))
    (loop for k from 1 to n
          for m = (cond ((null lst) n)
                        (t (pop lst)))
          do (princ m)
             (terpri))))
(main/2)Code language: Lisp (lisp)

無事に全問正解になりました。

3. 逆に考えると(AC18 418ms)

4. 出力回数を減らす(60ms)

418msのほとんどは、リスト生成と出力です。
同様のロジックでも、これを省略すると高速化できます。

(defun main/3 ()
  (let* ((n (read))
         (s (read-line))
         (out (make-string (* n 8) :element-type 'base-char))
         (w 0)
         (cnt 0))
    (labels ((put-char (ch)
               (setf (char out w) ch)
               (incf w))
             (put-digits (x)
               ;; 上の桁から順に out へ置く
               (multiple-value-bind (q r) (floor x 10)
                 (unless (zerop q) (put-digits q))
                 (put-char (code-char (+ 48 r)))))
             (emit (x)
               (declare (type fixnum x))
               (put-digits x)
               (put-char #\Newline)))
      (loop for idx below n
            for ch across s
            when (char= ch #\x)
              do (emit (1+ idx))
                 (incf cnt))
      (loop repeat (- n cnt)
            do (emit n)))
    (write-string out *standard-output* :end w)))

(main/3)
Code language: Lisp (lisp)

回答の出力に、princの代わりに補助関数 emit を使って文字列バッファ out に書き込んで、最後にまとめて出力します。

4. 出力回数を減らす(60ms)