AtCoder Beginner Contest 051 B - Sum of Three Integers

お題箱より。

B - Sum of Three Integers

計算量の考え方に入門するのにとてもいい問題です。 $O(K)$ 解法について知りたい…というリクエストだったので、色々な計算量で解く方法を解説してみました。

$O(K^{3})$ 解法

全ての基本、全探索です。 $0 \le X, Y, Z \le K$ を満たす $(X, Y, Z)$ を全て試します。和が $S$ になれば、条件を満たす組み合わせが1つ見つかったとして答えに1を足します。

この解法の計算量は $O(K^{3})$ です。今回の問題では $K \le 2500$ まで大きくなりうるので、単純に3乗すると $1.5625 \times 10^{10}$ です。流石に間に合いませんね。

TLEコード

Submission #6666763 - AtCoder Beginner Contest 051

$O(K^{2})$ 解法

計算量を落とします。計算量を落とす王道の考え方の1つは、「変数のうちいくつかを固定すると、残りが自動的に決まる」という性質がないか探してみることです。

今回の問題では、 $X, Y$ を固定すると残りの $Z$ は $S - (X+Y)$ として唯一に決まります。そのため $(X, Y)$ のペアについては全探索し、 $Z = S - (X+Y)$ が $0 \le Z \le K$ の範囲に入っていれば答えに1を足す、という解法で答えを求めることができます。

計算量は $O(K^{2})$ になります。今回の制約であればこれで十分間に合います。

ACコード

Submission #6666779 - AtCoder Beginner Contest 051

$O(K)$ 解法

固定する変数を3つ、2つ…と来たら次は1つです。 $X$ を固定したとしましょう。

このとき残りの $Y, Z$ については、 $Y+Z = S - X$ が成り立つことは分かるものの1通りには決まりません。では、この $(Y, Z)$ のペアの個数を求めることはできないでしょうか。

具体的な例として、入力が $K = 6, S = 11$ だとして、 $X = 3$ に固定したとします。このとき $0 \le Y, Z \le K$ かつ $Y+Z = S - X$ を満たす $(Y, Z)$ のペアは $(2, 6), (3, 5), (4, 4), (5, 3), (6, 2)$ の5通りです。このように $Y, Z$ それぞれのとり得る値はある連続した整数になっていて、そのとり得る整数の個数がそのままペア $(Y, Z)$ の個数になります。この範囲の両端を求めれば個数が分かりそうです。

満たすべき条件は $0 \le Y, Z \le K$ と $Y + Z = S - X$ なので、変数 $Z$ を消去すると以下の不等式が得られます。

$0 \le Y \le K$
- 変形して、 $Y \le S - X$ かつ $S - X - K \le Y$

つまり $Y$ の下限は $Y_{\min} = \max(0, S - X - K)$ 、上限は $Y_{\max} = \min(K, S - X)$ となります。もし $Y_{\min} \le Y_{\max}$ であれば範囲内の値は $Y_{\max} - Y_{\min} + 1$ 通りあり、そうでない場合は存在しません。より正確に書くと $\max(0, Y_{\max} - Y_{\min} + 1)$ 通りです。

$X$ を固定すればこの値は計算できるので、全探索した $X$ それぞれに対してこの値を求めて合計すれば解くことができます。これで $O(K)$ です。

ACコード

Submission #6666798 - AtCoder Beginner Contest 051

$O(1)$ 解法

$O(K)$ で終わりかと思いきや、実は $O(1)$ でも解けます。これは数学の力を使います。

先ほどと同様に入力が $K = 6, S = 11$ のときにおいて、 $X$ を全探索して $0, 1, 2, ...$ と1ずつ増やしていくときに、 $Y_{\min} = \max(0, S - X - K)$ 、 $Y_{\max} = \min(K, S - X)$ がどのように変化するかをプロットすると次のようになります。これは折れ線グラフであり、つまり傾きが同じそれぞれの部分は等差数列と見なすことができます。