【AtCoder ABC455D】
カード移動とクエリの簡約化

1. 問題

カード $1,\dots,N$ と山 $1,\dots,N$ があり、初期状態では山 $i$ にカード $i$ だけが積まれています。

$Q$ 回の操作で、カード $C_i$ とその上にあるカードすべてを、順序を保ったままカード $P_i$ の上へ移動します。

各操作の直前には、$C_i$ と $P_i$ は異なる山にあり、$P_i$ はその山の一番上にあることが保証されています。

すべての操作後、各山 $i$ に残っているカード枚数 $A_i$ を求め、$A_1,A_2,\dots,A_N$ を空白区切りで出力します。

制約は $1 \leq N,Q \leq 3 \times 10^5$ です。

D – Card Pile Query

2. スッキリで素朴に解いた（TLE14）

この山は、上に積み上げていくので「スタック」構造で管理できます。

• 【Common Lisp】コンスセルは部分リストではなくスロットでもある（リストとpop） – Chiilabo Note

(defun make-piles (n)
  (loop for i from 1 to n
        collect (list i)))
;; (make-piles 5) ;=> ((1) (2) (3) (4) (5))


(defun take-cards-in-piles! (piles card)
  (loop for cell on piles
        for pos = (position card (car cell))
        when pos
          return (loop repeat (1+ pos)
                       collect (pop (car cell)))))
;; (take-cards-in-piles! '((1) (2) (3) (4) (5)) 5) ;=> (5)

(defun move-cards-to-pile (piles c p)
  (let ((cards (nreverse (take-cards-in-piles! piles c))))
    (loop for cell on piles
          for top = (and (car cell) (caar cell))
          when (eql top p)
            do (loop for card in cards
                     do (push card (car cell)))
               (return)
          )
    piles))
;; (move-cards-to-pile '((1) (2) (3) (4) (5)) 5 4) 
;=> ((1) (2) (3) (5 4) NIL)


(defun main ()
  (let* ((n (read))
         (q (read))
         (piles (make-piles n)))
    (loop repeat q
          for cn = (read)
          for pn = (read)
          do (move-cards-to-pile piles cn pn))
    (loop for pile in piles
          do (format t "~a " (length pile)))))
#-swank(main)
Code language: Lisp (lisp)

3. 最後に何の上にあるかだけ（312 ms）

ここで思ったのが、逆順にたどることです。

1 3
4 5
1 4
4 2

たとえば、「4」のカードは、途中で 5 の上に来ましたが、最終的は 2 の上にあります。
また、4が山の一番上に乗っている間に、乗っかっても最終的に4の上に何が乗っかっているかは、この操作とは関係ありません。
つまり、4 5はなくても結果は同じです。

つまり、クエリは

1 4
4 2

に簡約化できます。

まずは、クエリのペアの c で昇順にstable-sortで並べ替えて、同じ c が続くときはスキップし、最後のクエリだけを取るようにしました。

(defun read-query-pairs (q)
  (loop repeat q
        collect (cons (read) (read))))
;;=> ((1 . 3) (4 . 5) (1 . 4) (4 . 2))

(defun simplificate-queries (qs)
  (loop for rest on (stable-sort (copy-seq qs)
                                 (lambda (a b) (< (car a) (car b))))
        while (cdr rest)
        for a = (caar rest)
        for b = (caadr rest)
        when (not (= a b))
          collect (car rest) into lst
        finally (return (append lst (list (car rest))))))
;; (simplificate-queries '((3 . 1) (5 . 4) (2 . 5) (5 . 7) (2 . 3) (6 . 2) (3 . 4) (5 . 1)))
;;=> ((2 . 3) (3 . 4) (5 . 1) (6 . 2))Code language: Lisp (lisp)

ただし、このままクエリにある数字 cを pの山に加算していってもうまくいきませんでした。

((2 . 3) (3 . 4) (5 . 1) (6 . 2))

expected: 2 0 0 4 0 0 1 
failed  : 2 1 0 3 0 0 1 Code language: Lisp (lisp)

たとえば、6 を 2 の上に移動するときに、2はすでに４の山に移動している、ということがあるからです。

そこで、rec-toでは、再帰的にハッシュテーブルを辿って、最終的な移動先を求めるようにしました。

(defun sum-assemble-queries (n query-pairs)
  (let* ((counts (make-array n :element-type 'fixnum
                               :initial-element 1))
         (selected-queries (simplificate-queries query-pairs))
         (map-table (make-hash-table)))
    (labels ((rec-to (from)
               (let ((memo (gethash from map-table)))
                 (setf (gethash from map-table)
                       (cond
                         ((null memo) from)
                         ((= memo from) from)
                         (t (rec-to memo)))))))
      (loop for (from . to) in selected-queries
            do (setf (gethash from map-table) to))
      (loop for (from . _) in selected-queries
            for final-to = (rec-to from)
            do (incf (aref counts (1- final-to)) 1)
               (setf (aref counts (1- from)) 0)))
    counts))
;;  (sum-assemble-queries 5 qs) 
;; #(0 3 1 0 1)Code language: Lisp (lisp)

これで、無事に答えが出ました。

3.1. 完成コード

(defun read-query-pairs (q)
  (loop repeat q
        collect (cons (read) (read))))
;;=> ((1 . 3) (4 . 5) (1 . 4) (4 . 2))
;; (defparameter qs '((1 . 3) (4 . 5) (1 . 4) (4 . 2)))

(defun simplificate-queries (qs)
  (loop for rest on (stable-sort (copy-seq qs)
                                 (lambda (a b) (< (car a) (car b))))
        while (cdr rest)
        for a = (caar rest)
        for b = (caadr rest)
        when (not (= a b))
          collect (car rest) into lst
        finally (return (append lst (list (car rest))))))
;; (simplificate-queries '((3 . 1) (5 . 4) (2 . 5) (5 . 7) (2 . 3) (6 . 2) (3 . 4) (5 . 1)))
;;=> ((2 . 3) (3 . 4) (5 . 1) (6 . 2))

(defun sum-assemble-queries (n query-pairs)
  (let* ((counts (make-array n :element-type 'fixnum
                               :initial-element 1))
         (selected-queries (simplificate-queries query-pairs))
         (map-table (make-hash-table)))
    (labels ((rec-to (from)
               (let ((memo (gethash from map-table)))
                 (setf (gethash from map-table)
                       (cond
                         ((null memo) from)
                         ((= memo from) from)
                         (t (rec-to memo)))))))
      (loop for (from . to) in selected-queries
            do (setf (gethash from map-table) to))
      (loop for (from . _) in selected-queries
            for final-to = (rec-to from)
            do (incf (aref counts (1- final-to)) 1)
               (setf (aref counts (1- from)) 0)))
    counts))
;; (sum-assemble-queries 5 qs) 
;;=> #(0 3 1 0 1)

(defun main/simplificate ()
  (let* ((n (read))
         (q (read))
         (qs (read-query-pairs q)))
    (loop for c across (sum-assemble-queries n qs)
          do (format t "~a " c))))
#-swank(main/simplificate)Code language: Lisp (lisp)