AGC 013 E 問題の解説

AGC013E.txt

＊はじめに
本記事は AtCoder Grand Contest 013 の E 問題の解説です。

問題文の閲覧やコードの提出はこちらから。
E: Placing Squares - AtCoder Grand Contest 013 | AtCoder
http://agc013.contest.atcoder.jp/tasks/agc013_e


＊執筆経緯
公式の問題解説 PDF や解説放送で紹介されていたこの問題の想定解法が、
個人的にはそれなりに高度な発想力を要する奇抜な解法だったと感じたので、
ここではもう少し思い付けそうな感じのする王道的な解法を紹介したいと思います。


＊問題概要
M 箇所に印の付いた長さ N の棒がある（棒の長さと印の位置は全て整数）。
印の付いた箇所で分割してしまわないように棒をいくつかの正整数長の小区間に分割する。
このとき、この分割のスコアを「各小区間の長さの２乗の総積（総和ではない）」と定める。
全ての可能な分割に対し、それらのスコアの総和を 10^9+7 で割った余りを求めよ。
  1 <= N <= 10^9
  0 <= M <= 10^5


＊解説
まず、可能な分割方法は全部で 2^(N-1-M) 通りもあるので、素朴に全探索してしまうと
計算量が O(N*2^N) になってしまいどうにもなりません（それはそう）。
そんなわけで、この問題は DP（動的計画法）で解きます。
ある程度競プロに慣れている方なら、次のような DP を思い付けるのではないでしょうか。

  dp[i] = 端から距離 i までの部分の分割に対するスコアの総和
        = 棒の長さが i だったとした場合のこの問題の答え

  初期値: dp[i] = i^2 （あるいは、dp[0] = 1, それ以外 = 0 としてもよい）

  i から j への遷移（i < j）
  ・i に印なし: dp[j] += dp[i] * (i-j)^2
  ・i に印あり: なにもしない

  答え: dp[N]

この DP はこのままでは計算量が O(N^2) なので実行制限時間に間に合いませんが、
実は遷移の計算方法を工夫することで高速化できます。
この DP を配る DP として考えてみると、i からの遷移の仕方は

  dp[i+1] += dp[i] *  1
  dp[i+2] += dp[i] *  4
  dp[i+3] += dp[i] *  9
  dp[i+4] += dp[i] * 16
  ...

という風になるので、DP テーブルを関数としてみた場合、i からの遷移とは、
DP テーブルに対して２次多項式（２次関数）を加算するという操作に他なりません。
すると、初期状態の DP テーブルの内容は２次多項式（dp[i] = i^2）であり、
また２次多項式に２次多項式を加算しても（高々）２次多項式であることから、
DP テーブルの内容は常に２次多項式として表すことができます。
さらに、２次多項式には、「３点での値がわかっていれば元の２次多項式を復元できる」
という性質があります（ついでに言うと、n 次多項式なら n+1 点あれば復元できます）。
というわけで、DP テーブルの全てを持っておく必要はなく、
例えば dp[i], dp[i+1], dp[i+2] の３点だけ覚えておけばいいということがわかります。

この工夫により、計算量を O(N) にできることがわかりましたが、
N が最大ケースで 10^9 であることを考えるとまだ厳しいです。
遷移の計算にはまだ高速化の余地があるので、この部分をもう少し詳細に見てみましょう。

覚えておく３点として例えば dp[i], dp[i+1], dp[i+2] を選んだ場合、これらの値から
dp'[i+1], dp'[i+2], dp'[i+3] を計算することになります（dp'[x] は更新後の DP テーブル）。
このとき dp[i+3] の値が必要になるので、dp[i], dp[i+1], dp[i+2] の値から
２次多項式を復元しなければなりません。
これは３式からなる連立方程式を解くことでも求められますが、
この３点に対しラグランジュ補間をすれば次のように一発で求まります。

           x-(i+1)   x-(i+2)                 x-i       x-(i+2)                   x-i       x-(i+1)        
  f(x)  =  ------- * ------- * dp[i]  +  ------- * ----------- * dp[i+1]  +  ------- * ----------- dp[i+2]
           i-(i+1)   i-(i+2)             (i+1)-i   (i+1)-(i+2)               (i+2)-i   (i+2)-(i+1)        

これに i+3 を代入すれば、dp[i+3] の値は

  dp[i+3] = f(i+3) = dp[i] - 3*dp[i+1] + 3*dp[i+2]

であるとわかります。
そして、i の位置に印がある場合はこれでおしまい、印がない場合は i からの遷移として
dp[i+1], dp[i+2], dp[i+3] にそれぞれ dp[i], 4*dp[i], 9*dp[i] を加算するので、
まとめると次のようになります。

・印あり
  dp'[i+1] = dp[i+1]
  dp'[i+2] = dp[i+2]
  dp'[i+3] = dp[i] - 3*dp[i+1] + 3*dp[i+2]

・印なし
  dp'[i+1] = dp[i+1] + dp[i]
  dp'[i+2] = dp[i+2] + 4*dp[i]
  dp'[i+3] = 10*dp[i] - 3*dp[i+1] + 3*dp[i+2]

これは行列を用いて表すと、次のようになります。

・印あり
  dp' = B * dp

・印なし
  dp' = A * dp

        ( dp'[i+1] )       ( dp[i]   )
  dp' = ( dp'[i+2] )  dp = ( dp[i+1] )
        ( dp'[i+3] ),      ( dp[i+2] )

      (  1  1  0 )      ( 0  1  0 )
  A = (  4  0  1 )  B = ( 0  0  1 )
      ( 10 -3  3 ),     ( 1 -3  3 )

ここで列ベクトル dp の初期値は (1, 4, 9)（を転置したもの）であり、
これに各 i での印の有無に応じて遷移行列 A もしくは B を左から順番に掛けていくことで答えが求まるのですが、
行列の積では結合律が成り立つため、先に遷移行列の積をまとめて計算してしまってから
列ベクトルに掛けるという順番で計算しても同じ答えが得られます。
最大ケースでは印のない箇所の方が圧倒的に多いことから大部分が A の積の計算となるので、
A の累乗にバイナリ法（繰り返し二乗法、ダブリング）を用いて高速化すれば、
行列の積の計算回数のオーダーが O(M*log(N)) となり、実行制限時間に間に合います。

なお、実際に答えるのは mod 10^9+7 での値ですが、これは単に計算過程で適宜 mod をとるだけでよいです。
ただし、負数の剰余が負になるような言語では、負数の取り扱いに注意しましょう。


＊別解
もしかするとこちらの解法の方が思い付きやすいかも知れません。

先程の解法では、DP テーブルの内容を保持するために dp[i], dp[i+1], dp[i+2] を
記憶していましたが、その代わりに値と差分と２階差分、つまり

  a := dp[i]
  b := dp[i+1] - dp[i]
  c := (dp[i+2] - dp[i+1]) - (dp[i+1] - dp[i])

を記憶しておくという方法もあります。
２次多項式は２階差分をとると定数になり、３階差分をとると恒等的に 0 となって消えてしまうので、
この３つの値を記憶しておけば元の２次多項式を復元できるという理屈です。

まず、初期状態の DP テーブルは〈 1, 4, 9, 16, ... 〉となっているので、
これは (a, b, c) = (1, 3, 2) と表現されます。
そして遷移について考えると、まず印がある場合は

  (a, b, c) |----> (a+b, b+c, c)

のように更新します。そして、印がない場合もまずは同様にし、
その後で DP テーブル（の i+1 番目以降）に〈 a, 4*a, 9*a, 16*a, ... 〉を加算するのですが、
これは (a, 3*a, 2*a) で表される関数を加算するということであり、
また、差分の線型性からこれは単に (a+b, b+c, c) との和となるので、

  (a, b, c) |----> (2*a + b, 3*a + b + c, 2*a + c)

となります。行列として書き直すと

      ( 2 1 0 )      ( 1 1 0 )
  A = ( 3 1 1 )  B = ( 0 1 1 )
      ( 2 0 1 ),     ( 0 0 1 )

となります。

こちらの解法の方が遷移行列の導出が若干楽でしょうか。