ABC138D - Kiを解いてみた

40代文系社内SEが、AtCoder Beginner Contest 138 D - Kiを解いてみました。

まず問題文を確認します。

根が頂点1の木があり、それぞれの頂点はカウンターをもっています。部分木のカウンターに値を加算する操作を繰り返し、最終的なカウンターの値を求めます。

まず、「木」と記載されていることからグラフの問題であることが考えられます。また、部分木に値を加算するとありますが、おそらくそのまま実装するとTLEになるだろうと思われます。

次に制約を確認します。

頂点の個数$N$が最大$2 \times 10 ^ 5$であり、操作回数$Q$も最大$2 \times 10 ^ 5$であることから、毎回頂点1に操作をするケースで計算量が$10 ^ {10}$を超えて間に合いません。

どのように計算量を抑えるかと考えたとき、根から部分木に値を加算していくという操作が累積和のようにみえました。すべての頂点が一直線につながっているようなケースを考えると、イメージしやすいかと思います。

累積和が使えれば、値を加算する操作は最後に1回だけ処理すればいいので、計算量を$2 \times 10 ^ 5$に抑えられそうです。

以上より、以下の手順で実装します。

1. a、bを隣接リストとして読み込む
2. カウンターを用意して、カウンターpにxを加算する
3. 頂点1からグラフを探索し、累積和の要領で値を加算する

(defun solve ()
  (let* ((n (read))
         (q (read))
         (graph (make-array (1+ n) :initial-element nil))
         (counter (make-array (1+ n) :initial-element 0)))

    ;; a, bを読み込み、隣接リストとして保持
    (loop repeat (1- n)
          for ai = (read)
          for bi = (read)
          do (push bi (aref graph ai))
             (push ai (aref graph bi)))

    ;; p, xを読み込み、カウンターに記録
    (loop for i from 1 to q
          for p-i = (read)
          for xi  = (read)
          do (incf (aref counter p-i) xi))
    (let ((visited (make-array (1+ n) :initial-element nil)))
      ;; 頂点1を訪問済にして深さ優先探索を開始
      (setf (aref visited 1) t)
      (dfs graph counter '(1) visited)
      (format t "~{~A~^ ~}~%" (coerce (subseq counter 1) 'list)))))

(defun dfs (graph counter visiting visited)
  ;; 深さ優先探索でグラフを探索し、カウンターを更新する
  (when (not (null visiting))
    (loop with pos = (pop visiting)
          for dest in (aref graph pos)
          when (not (aref visited dest))
            do (incf (aref counter dest) (aref counter pos)) ; 次の頂点にカウンターの値を加算する
               (setf (aref visited dest) t)
               (push dest visiting))
    (dfs graph counter visiting visited)))

(solve)

ACできたのですが、訪問済のチェックを省略していて1回WAになってしまいました。