全探索（Sum of Three Integers）
はじめに
AtCoderのB問題に「Sum of Three Integers」という全探索をいかに効率的にできるかが本質の問題があり、僕のようなアルゴリズム初心者にはすごく面白い内容だったので紹介したいと思う。
問題内容
二つの整数K,Sと三つの変数X,Y,Zがある。
X,Y,Zは、「0 <= X,Y,Z <= K」であるとき、「X + Y + Z = S」を満たすX,Y,Zの割り当ては何通りであるか。
制約として、「2 <= K <= 2500」、「0 <= S <= 3K」が与えられる。
解法
おそらく、一番最初に思いつくことは3重のfor文だと思う。XとYとZをそれぞれ、０からKまで回し、足し算の結果がSのとき、カウントに1を足していく。
しかし、これでは、Kが99の時点でそれぞれで100の3乗となり、1万回の計算が必要となってくる。もし、Kに2500が割り当てられるとプログラムがタイムアウトしてしまいます。
この問題は、変数を減らすことが解くカギになっていると感じます。例えば、XとYを定数にすれば、Zは、S-(X+Y)となり、Zが条件を満たせば、カウントを増やせばよい。このプログラムでは、計算量がO(K^2)まで減らすことができる。これでタイムアウトを防ぐことはできた。以下がそのプログラムである。
public class Main {
public static void main(String args)
throws Exception {
Scanner scanner = new Scanner(System.in);
int K = scanner.nextInt();
int S = scanner.nextInt();
int sub = 0;
int count = 0;
for(int i=0;i<=K;i++){
if(i != S){
for(int j=0;j<=K;j++){
if(i+j != S){
sub = S - (i+j);
if(sub>=0 && sub<=K){
count++;
}
}else if(i+j == S)
{
count++;
break;
}
}
}else if(i == S){
// System.out.println("i:"+i);
count++;
}
}
System.out.println(count);
}
}
for文の中で使っているif文はiがSの時に無駄に内側のfor文に入らないようにするためである。
このプログラムでも十分、正解になるのだが、一応計算量がさらに減り、O(K)になる方法もある。それは、変数を二つ減らす方法である。
僕のプログラムでは、Xを定数でfor文を回すことで実現した。
Xを固定することで、「Y+Z = S-X」という等式が成り立つ。次にもう一つ、変数を無くしたいのでZをなくなるように調整する。Zは「S-X-Y」であるため、「0 <= S-X-Y <= K」が成り立つ。この不等式から「Y<=S-X」「S-X-K<=Y」という条件が出てくる。これは、Xを固定したときに、Yは最大でどのくらいの数値がでてくるのか、逆に最低でどのくらいの数値が出てくるのかが分かる。これがわかれば、Yは連続した値で出てくるはずなので（最大値ー最低値＋１）がXがある定数のときの条件を満たすものとなる。
プログラムは以下となる。
public class Main {
public static void main(String args)
throws Exception {
Scanner scanner = new Scanner(System.in);
int K = scanner.nextInt();
int S = scanner.nextInt();
int Y_max = 0;
int Y_min = 0;
int count = 0;
for(int x=0;x<=K;x++){
Y_min = S - x - K;
Y_max = S - x;
if(Y_min<0)Y_min = 0;
if(Y_max>K)Y_max = K;
if(Y_max>=Y_min)count += Y_max - Y_min + 1;
System.out.println("Y_max: "+Y_max+" Y_min: "+Y_min);
}
System.out.println(count);
}
}
条件は先ほど書いたが、この条件の前にYには０以上K以下というものがある。
なので、Yの最大値はKより大きくならないように、最小値は０より小さくならないように条件をつけている。これで、計算量はO(K)となる。
以上が、「Sum of Three Integers」の二つの解法である。
