「Xmas Contest 2015 昼の部」のA問題を丁寧に解いてみる

# 概要
　発端、[AtCoder](http://atcoder.jp/)代表取締役社長の[chokudai](https://twitter.com/chokudai)氏と、新卒文系プログラマ兼マンガ家｜腐女子の[ちょまど](https://twitter.com/chomado)氏との次のようなやり取りからでした。
　　[https://twitter.com/chokudai/status/690039649268580352](https://twitter.com/chokudai/status/690039649268580352)
　端的に言えば、「A問題が解ければプログラマなんですよね？**じゃあこれ解いてください＞＜**」というchokudai氏からの挑戦状ですね。あまりの難易度にちょまど氏は早々にギブアップし、後から参加した私も**その難易度に目を剥きました**。
　以下はこの問題の解説記事となります。ネタバレを避けたい人は何も見ずにまずは解いてみましょう＞＜
　　[A - Accumulation](http://xmascontest2015noon.contest.atcoder.jp/tasks/xmascontest2015_a)

# まずは問題の意図をつかみましょう
　問題文を読めば分かりますが、随分と**数学っぽい**問題です。それも、整数論の香りがビンビンにします。
　幸い擬似コードが与えられているので、「どんな処理をしているか」は簡単につかむことができますね。Rubyで書くとこんな感じです。

```code1.rb
=begin
入力：
  N
  X,T,A,B,C
入力の条件：
  1 ≦ N ≦ 10^6
  1 ≦ T,C ≦ 10^9
  0 ≦ X,A,B ＜ C
=end

N = gets.to_i
x,T,A,B,C = gets.split.map(&:to_i)
s = 0
N.times{
  s += x
  T.times{
    x = (A * x + B) % C
  }
}
puts s
```

　……こうして見ると分かるように、素朴なコードだと「x = (A * x + B) % C」がN×T回実行されることが分かります。入力の条件から、それが最大10^15(100京)回も繰り返されることになるので、問題の条件(2sec)を考慮すると到底間に合いません。実際、提出してみると惨憺たる成績でした。
　　[http://xmascontest2015noon.contest.atcoder.jp/submissions/616130](http://xmascontest2015noon.contest.atcoder.jp/submissions/616130)

# 次に高速化を検討します
　上記ソースをよく読んでみると、「sにxを足す→xに対して処理する→sにxを足す→……」といった処理手順になっていることが分かります。この「xに対して処理する」部分に大きく計算時間を取られているのが問題なのです。うまく工夫して計算量を削れないでしょうか？
　そんな時に登場するのが整数論です。そこそこ長くなるので別に分けます。

## 整数論の基礎(ざっくり)
　まず、剰余計算について整理しましょう。整数Aを整数Bで割った余りが整数Cであるとき、適当な整数Nを用いて$A=B×N+C$と表現されます。小さな数で例えれば、A=25,B=7,C=4である際、$25=7×3+4$ということですね。
　ここで記号$\pmod{p}$を導入します。意味としては、**「法pの世界において」**となります。この世界の元では、**「pで割った際の余りが同じ数は等しいと見なす」**ようになります。小さな数で例えば、$25=18=11=4\pmod{7}$ということですね。
　この不思議な世界の元でも、四則演算は同様に成り立ちます(ただし除法については多少ややこしい)。つまり、4×6も8×3も24と等しくなります。ただし先ほどのルールから、$7×3+3=24=3=66=7×9+3\pmod{7}$となりますが。

## 計算量削減法その1：計算をまとめる
　この整数論の性質を用いると、「xに対して処理する」部分をコンパクトにまとめることが出来ます。
　法Cの世界の元でも四則演算が成り立つので、$f(x)=A×x+B$として、$g(x,t)=f(f(f(...(x))))$と置くと(ただしfがt個入れ子になっている)、適当な整数P,Qを用いて$g(x,t)=P×x+Q$と書くことができます。つまり、事前にPとQを計算することによって、「xに対して処理する」部分の計算が瞬時に終わることになるのです。
　PとQの計算方法ですが、$A×(P×x+Q)+B=A×P×x+A×Q+B$なので、$P←A×P, Q←A×Q+B$とする処理を(T-1)回繰り返せばオーケーです。これをRubyで書くと次のようになります。

```code2.rb
N = gets.to_i
x,T,A,B,C = gets.split.map(&:to_i)
p,q = A,B
(T - 1).times{
  p = (A * p) % C
  q = (A * q + B) % C
}
s = 0
N.times{
  s += x
  x = (p * x + q) % C
}
puts s
```

　ループ内でさり気なく、毎回Cで剰余を取っていることに注意してください。別にTループが終わった後に剰余を取ればいいのですが、pやqの桁数が大きくなりすぎるのを剰余によって抑えています。
　ただし、実のところ、これでもまだ計算量は多いです。もっと工夫しなければ、テストケースを全て通過させることはできません。
　　[http://xmascontest2015noon.contest.atcoder.jp/submissions/617963](http://xmascontest2015noon.contest.atcoder.jp/submissions/617963)

## 計算量削減法その2：バイナリ法
　そもそも、先に書いた通り、「f(x)をt回適用すること」は、「f(x)をt/2回適用してから、f(x)をt/2回適用する」ことと同じです(入れ子なので)。要するに、$g(x,t/2)=P'×x+Q'$と置くと、$g(x,t)=P'×(P'×x+Q')+Q'=P'×P'×x+P'×Q'+Q'$となります。
　このことから、問題サイズを1/2倍にしてからその結果を再利用することで、**計算量を大幅に削減できる**と分かります。この発想は一般に**バイナリ法**と呼ばれ、広くは**分割統治法の一つ**として分類されます。
　まずはこれを再帰関数として実装してみましょう。具体的はこんな感じです。

```code3.rb
def calc(a,b,c,t)
	return a,b if t == 1
	p1,q1 = calc(a,b,c,t/2)
	p,q = (p1*p1)%c,(p1*q1+q1)%c
	p,q = (a*p)%c,(a*q+b)%c if t % 2 == 1
	return p,q
end

N = gets.to_i
x,T,A,B,C = gets.split.map(&:to_i)
p,q = calc(A,B,C,T)
s = 0
N.times{
  s += x
  x = (p * x + q) % C
}
puts s
```

　ここでポイントとしては、tが偶数とは限らないことから、それを考慮して再帰関数内が冗長になっていることです。とはいえこれでも十分に速く、テストケースを楽々と通過します。
　　[http://xmascontest2015noon.contest.atcoder.jp/submissions/617992](http://xmascontest2015noon.contest.atcoder.jp/submissions/617992)
　より高速化したい場合は、再帰関数ではなく直接書き下せばよいでしょう。つまり、「f(x)を2^n回適用」した際の係数p,qを別途計算しておき、後はそれを繋ぎ合わせるだけで計算が完了します。

```code4.rb
def calc(p,q,a,b,c)
	return (a*p)%c,(a*q+b)%c
end
n = gets.to_i
x,t,A,B,C = gets.split.map(&:to_i)
p,q = 1,0
p_,q_ = A,B
while t >= 1
	if t % 2 == 1 then
		p,q = calc(p,q,p_,q_,C)
	end
	p_,q_ = calc(p_,q_,p_,q_,C)
	t /= 2
end
s = 0
n.times{
  s += x
  x = (p * x + q) % C
}
puts s
```

　実際、実行結果を見ても、再帰関数より速い結果になっていることが分かります。では、このコードが最適なのでしょうか？
　　[http://xmascontest2015noon.contest.atcoder.jp/submissions/618027](http://xmascontest2015noon.contest.atcoder.jp/submissions/618027)

## 計算量削減法その3：更に一工夫！
　ここでf(x)の式を思い出してみてください。あれを1回入れ子にすると$f(f(x))=A×A×x+A×B+B$となりますが、このまま入れ子にし続けるとどうなるのでしょうか？
　……答えは、$g(x,t)=A^t×x+A^{t-1}×B+A^{t-2}×B+...+A×B+B$です。A=0,1なら自明ですので、以下はA≧2で考えます。
　等比数列の和の公式から、$A^{t-1}+A^{t-2}+...+A+1=(A^t-1)/(A-1)$となりますので、$g(x,t)=A^t×x+(A^t-1)/(A-1)×B$となります。ただ、pがqで割り切れる際、$(p/q)\bmod r=(p \bmod qr)/q$が成立しますので、$(A^t-1)/(A-1) \bmod C=([A^t-1] \bmod [C(A-1)]) / [A-1]$と変形できます。以上を踏まえてソースに落としこむとこんな感じ。

```code5.rb
def powmod(a,t,p)
	r = 1
	b = a
	while t > 0
		if t % 2 == 1 then
			r = r * b % p
		end
		b = b * b % p
		t /= 2
	end
	return r
end
 
n = gets.to_i
x,T,A,B,C = gets.split.map(&:to_i)
 
p = powmod(A,T,C)
q = (powmod(A,T,C*(A-1))+(C*(A-1)-1))/(A-1)*B
 
s = 0
n.times{
  s += x
  x = (p * x + q) % C
}
puts s
```

　……え、この方が冗長だろって、そ、ソンナコトナイヨ。
　　[http://xmascontest2015noon.contest.atcoder.jp/submissions/618063](http://xmascontest2015noon.contest.atcoder.jp/submissions/618063)
　ちなみに[上記の短縮アイディアはこのツイートから](https://twitter.com/angel_p_57/status/690599136903593984)でした。提案者による、より洗練されたコードもご覧ください。
　　[http://xmascontest2015noon.contest.atcoder.jp/submissions/615387](http://xmascontest2015noon.contest.atcoder.jp/submissions/615387)

# まとめ
　途中から高速化というより知恵比べな方向に脱線しましたが、**「分割できれば楽になる」**ことは分かっていただけたでしょうか。
　ちなみに私の場合、最初は**[メモ化してなんとか乗り切った記憶があります＞＜](http://xmascontest2015noon.contest.atcoder.jp/submissions/615269)**