【AtCoder ABC432C】
Candy Tribulationの効率解
（合同式）

by
  • AtCoder ABC432C「Candy Tribulation」を合同式で解く方法で、計算量がO(N)となり、実行時間が1645msから59msに短縮されました。
  • 各子供への飴の総重量Wを、大小の重さの差Dを法とした剰余で統一できるかを確認します。
  • 上限から剰余のずれを引くことで、条件を満たす最大のWを直接計算で求めます。

AtCoder Daily Trainingで解いたのですが、より効率的な答えがあるので、考えてみます。

関連記事

1. 問題（Candy Tribulation）

X グラムの小さい飴と Y グラムの大きい飴（X < Y）を N 人の子供に配ります。
子供 i にはちょうど A[i] 個配り、全員の総重量が等しくなる条件で、大きい飴の合計個数の最大値を求めます。
配り方が存在しない場合は -1 を出力します。

C – Candy Tribulation

1.1. Wの候補を探索する方法

一人の子供に配る飴の総重量を W とおきます（weight）。

最初に書いた解法では、可能な W を候補としてループ内で 1 個ずつ試して回答を探しました。

(defun candy-weight (x y a n)
  (+ (* y n)
     (* x (- a n))))

(defun can-tribute-candy-p (x y a weight)
  (multiple-value-bind (n reminder)
      (floor (- weight (* x a))
             (- y x))
    (and (zerop reminder)
         (>= n 0) )))

(defun all-heavy-candys (x y An weight)
  (loop for a in An
        sum (floor (- weight (* x a))
             (- y x))))

(defun solve/2 (x y An)
  (let* ((min (reduce #'min An))
         (max (reduce #'max An))
         (minimum-weight (* max x)))
    (loop for n from min downto 0
          for weight = (candy-weight x y min n)
          while (>= weight minimum-weight)
          when (loop for a in An
                     always (can-tribute-candy-p x y a weight))
            do (return (all-heavy-candys x y An weight))
          finally (return -1))))

(defun main/2 ()
  (let* ((N (read))
         (x (read))
         (y (read))
         (An (loop repeat N collect (read))))
    (princ (solve/2 x y An))))

#-swank(main/2)Code language: Lisp (lisp)

ただ、回答の最長実行時間は、1645 ms 。
計算量は、重さ候補のループと子供ごとのチェックのループで、O(N2)です。
ほかの人の回答を読んでいたら、11 〜 89 ms という高速な解法がいくつもあります。

2. 整数条件と剰余による絞り込み

まず、子供 i に大きい飴を n_i 個配るとします。

整数条件と剰余による絞り込み ① 式の変形 W = A[i]·X + n_i·D D = Y − X（重さの差） n_i = (W − A[i]·X) / D ② 合同条件 W ≡ A[i]·X (mod D) 全員で余りが一致 → 解あり Wmax の求め方 7 12 17 22 27 上限25 Wmax = 22 diff = (上限 mod D − target + D) mod D Wmax = 上限 − diff Wmax ≥ 下限 → 解が存在

すると、小さい飴は A[i] – n_i 個配ることになります。
すると、それぞれの子供に配る飴の総重量を W は、以下の式になり、取りうる範囲はすべて小さな飴を配る〜すべて大きな飴を配る、の間にあります。

W = n_i * Y + (A[i] - n_i) * X

A[i] * X  ≤  W  ≤  A[i] * Y

次に、大小の重さの差 D = Y – X とおきます（difference）。
すると、総重量 W は次の関係式になります。

W  = A[i] * X + n_i * D

式を変形すると、
n_i = (W - A[i] * X) / D と表せます。

ここで、n_i が整数になるなら、

W ≡ A[i] * X  (mod D)Code language: CSS (css)

が成り立ちます。
これが全員で一致すれば、条件を満たす W が存在することになります。

2.1. 整数条件の確認（check-remainder）

一人目のDで割った余りを target = A[0] * X mod D とおきます。

すると、全員の余りが target と一致することを確認します。

(defun check-remainder (x d An)
  (loop for a in (cdr An)
        with target = (mod (* (car An) x) d)
        always (= (mod (* a x) d) target)))Code language: Lisp (lisp)

every で全員の余りが target と一致するか確認します。
1 人でもずれていれば nil を返します。

2.2. 最大の W を求める（max-valid-weight）

大きい飴の個数を最大化するには W を最大にすればよいです。
この W の最大値 Wmaxを求めます。

W は全員の範囲に収まるので、上限は min(A[i] * Y)、下限は max(A[i] * X) です。

たとえば、D = 5、target = 2 というときは、
W の候補は、2、7、12、17、22、27 と、D 刻みで並んでいます。

ここで、Wの上限 min_ay = 25 だったとします。
すると、この中の最大は Wmax = 22 ということになります。

これを数列を並べずに計算で求めるには、D を法とした減算を使います。
なぜなら、25より小さくて、5 の余りが 2 の数を計算します。

diff = (0 - 2 + 5) mod 5 = 3
Wmax = 25 - 3 = 22

一般化すると、min_ay の Wmax の差を D で割った余りは、
min_ay と target の差を D で割った余りと同じで、この差 diff から逆算すれば、Wmaxが計算できるのです。

diff = (min_ay mod D - target + D) mod D
Wmax = min_ay - diff

引く量は常に 0 以上 D 未満なので、Wmax は min_ay – (D-1) 以上です。
+ D は (rem - target) が負になったときに mod の結果が負のまま返るのを防ぐための調整です。

(defun max-valid-weight (x y d An target)
  (let* ((upper-bound (reduce #'min An
                              :key (lambda (a) (* a y))))
         (lower-bound (reduce #'max An
                              :key (lambda (a) (* a x))))
         (rem    (mod upper-bound d))
         (diff   (mod (- rem target) d))
         (wmax   (- upper-bound diff)))
    (cond
      ((>= wmax lower-bound) wmax)
      (t nil))))Code language: Lisp (lisp)

reduce #'min と reduce #'max に :key で A[i] の計算を渡すことで、リストを別途作らずに一度の走査で求めています。
Wmax が下限を下回れば nil を返します。

2.3. 大きい飴の合計（total-large）

Wmaxが求まれば、大きい飴の個数はかんたんに計算できます。

(defun total-large (x d An wmax)
  (loop for a in An
        sum (floor (- wmax (* a x)) d)))Code language: Lisp (lisp)

各人の ni = (Wmax – A[i] * X) / D を合計します。

2.3. 大きい飴の合計（total-large）

実行時間は、1645 ms から 59msになりました。

3. 全体

(defun check-remainder (x d An)
  (loop for a in (cdr An)
        with target = (mod (* (car An) x) d)
        always (= (mod (* a x) d) target)))

(defun max-valid-weight (x y d An target)
  (let* ((upper-bound (reduce #'min An
                              :key (lambda (a) (* a y))))
         (lower-bound (reduce #'max An
                              :key (lambda (a) (* a x))))
         (rem    (mod upper-bound d))
         (diff   (mod (- rem target) d))
         (wmax   (- upper-bound diff)))
    (cond
      ((>= wmax lower-bound) wmax)
      (t nil))))

(defun total-large (x d An wmax)
  (loop for a in An
        sum (floor (- wmax (* a x)) d)))

(defun solve/mod (x y An)
  (let* ((d (- y x))
         (target (mod (* (car An) x) d)))
    (unless (check-remainder x d An)
      (return-from solve/mod -1))
    (let ((wmax (max-valid-weight x y d An target)))
      (if wmax
          (total-large x d An wmax)
          -1))))

(defun main/mod ()
  (let* ((n (read))
         (x (read))
         (y (read))
         (An (loop repeat n collect (read))))
    (princ (solve/mod x y An))))

#-swank (main/mod)Code language: Lisp (lisp)

全体の計算量は O(N) です。