JUMP問題 - ユニファ開発者ブログ

こんにちは、プロダクト開発部のちょうです。前回から結構長い時間経ってましたので、すこし別のトピックにしようと思います。紹介なしで短い時間でできることの一つは多分アルゴリズム問題でしょう。今回こういう問題を解いてみたいと思います。

入力：整数の配列、例えば [3, 1, 0, 4]。あなたは最初に最初の要素の場所に居て、それとその要素の値によってあなたが最大どれぐらいの距離でJUMPできるかが決められる（要素の値はすべて0以上）。例えば [3, 1, 0, 4] の最初の要素が 3 なので、あなたが最大3ステップJUMPできる。もちろん、1ステップと2ステップもできる。いま、問題は入力された配列で最後の要素までJUMPできるかどうかを出力してください。出力は true と false のどちらでよい。例えば [3, 1, 0, 4] の最初の要素 3 から3ステップJUMPすれば最後の要素4までJUMPできる。また、[2, 1, 0, 4] のような配列では、最初の要素 2 からステップ1でもステップ2でも、次の要素から最後の要素までJUMPできない。

解答の関数は以下のようになる（言語はJava、入力と出力の仕方は考慮しなくてもいい）。

class Solution {
    public boolean canJump(int[] nums){
    }
}

この問題を見て最初に考えられるのは、いま居る位置の要素の値ですべての可能性を確認します。ここで再帰を使うと分かりやすいでしょう。

class Solution {
    public boolean canJump(int[] nums){
        int n = nums.length;
        if(n == 0) {
            return false;
        }
        if(n == 1) {
            return true;
        }
        return doCanJump(nums, n, 0);
    }

    private boolean doCanJump(int[] nums, int n, int i) {
        int availableSteps = nums[i];
        if(i + availableSteps >= n - 1) {
            return true;
        }
        for(int j = 1; j <= availableSteps; j++) {
            if(doCanJump(nums, n, i + j)) {
                return true;
            }
        }
        return false;
    }
}

関数 canJump では、まず配列をチェックします。空の場合はそもそもJUMPできないので、結果は false 。それと、要素数は1の場合、この唯一の要素の値はどうであれ、最初でも最後の要素なので、もちろんJUMPできます。次は再帰関数 doCanJump です。再帰関数 doCanJump はまずいまの要素の値とその位置をたしてみて、もし最後の要素かその以上の位置の要素に届けるなら、JUMPできるとなります。そうでない場合、1から最大まですべての可能性をテストします。すべての可能性でJUMPできるルートがないなら、JUMPできないと返します。

このコード理論上かならずJUMPできるルートを探し出します。ただ、すべてのルートをチェックしているので、最悪の場合、O(N^2) の時間複雑度でルートを探します。一応、コードを改善する方法はあります。例えば

ステップ1から最大までではなく、最大からステップ1まで試す
再帰をループにする

しかし、これらの改善策をしても計算量が変わりません。計算量を減らすには、終点（最後の要素）までの要素に「この要素なら終点までJUMPできる」ようなキャッシュを作る方法があります。例えば、配列 [3, 1, 2, 4, 6, 3, 1, 3] という配列で

最後から二番目の要素 1 から終点までJUMPできる
最後から三番目の要素 3 から終点までJUMPできる
最後から四番目の要素 6 から終点までJUMPできる
最後から五番目の要素 4 から終点までJUMPできる
最後から六番目の要素 2 から最後から五番目と四番目までJUMPできる。かつどっちも終点までJUMPできるので、最後六番目からも終点までJUMPできる
最後から七番目の要素 1 から最後から五番目までJUMPできる。かつ最後から五番目が終点までJUMPできるので、最後から七番目も終点までJUMPできる。
最後から八番目の要素 3 から最後から七、六、五番目までJUMPできる。かつどっちも終点までJUMPできるので、最後から八番目つまり最初の要素から終点までJUMPできる。

つまり、最初からではなく、最後からきっと使われるだろうと思う要素を事前に計算すれば、最初の要素からJUMPかどうかはわかります。こういう考えをベースにしたコードです。

class Solution {
    public boolean canJump(int[] nums){
        int n = nums.length;
        if(n == 0) {
            return false;
        }
        if(n == 1) {
            return true;
        }
        boolean[] dp = new boolean[n];
        dp[n - 1] = true;
        for(int i = n - 2; i >= 0; i--) {
            if(i + nums[i] >= n - 1) {
                dp[i] = true;
                continue;
            }
            for(int j = nums[i]; j >= 1; j--) {
                if(dp[i + j]) {
                    dp[i] = true;
                    break;
                }
            }
        }
        return dp[0];
    }
}

要素数が1以上の部分を見てみましょう。名前が dp という配列を作ります。これは「この要素から最後までJUMPできるか」のキャッシュです。もちろん、最後の要素はJUMPできるので、 true にします。次は最後から二番目から各要素を計算します。もし要素の位置と値から最後までかそれ以上かに届けるなら、JUMPできるとなります。もしくは、要素からJUMPできる範囲内にある要素がJUMPできるならこの要素からもJUMPできます。最後最初の要素の計算結果を返します。

計算結果をキャッシュすることで、計算量が回帰関数よりだいぶ減りましたでしょう。ちなみにですが、配列の名前は cache ではなく dp だった理由は、この方法が「Dynamic Programming」（動的計画法）と呼ばれるのです。動的計画法を知らなくても問題ありません。小さいな問題の結果を記録して大きな問題を解くのを覚えればいいです。

解答はここで終わっていません。すべの可能性を確認するので、理論上の計算量もそのまま O(N^2) です。すべての可能性を確認せずこの要素ならきっとJUMPできるような方法があるのでしょうか。

もう一度配列を見てみましょう。例えば [3, 0, 2, 2, 3, 2, 1, 3]。最初の要素 3 から1ステップから3ステップまでJUMPできますが、どれを選択すれば最後までJUMPできますか。すぐ隣の 0 を選択すればJUMPできなくなるのでだめです。次の２つ 2 どれもよさそうですが、一番目の 2 を選ぶより二番目の 2 がいいでしょう。なぜなら、一番目の 2 からJUMPできる位置も、二番目の 2 からもJUMPできます。逆に、二番目の 2 からJUMPできる位置は一番目の 2 からJUMPできない位置があります。もうすこしルール化にすると、

JUMPできる位置の中で、位置＋要素の値が一番大きい要素を選ぶ。もし複数いれば、一番右のを選ぶ。

実際やってみましょう。配列 [3, 0, 2, 2, 3, 2, 1, 3] 最初の要素 3 からJUMPできる範囲で、二番目の 2 を選ぶことになります。また、二番目の 2 から、隣の 3 と 2 から右の 2 を選びます。選んだ 2 から最後の要素までJUMPできます。これでチェックするルートは一つになり、計算は楽になるでしょう。

以上のルールは「Greedy Algorithm」（貪欲法）と似たようです。コードは以下です。

class Solution {
    public boolean canJump(int[] nums){
        int n = nums.length;
        if(n == 0) {
            return false;
        }
        if(n == 1) {
            return true;
        }
        int current = 0;
        while (true) {
            int step = nums[current];
            if (current + step >= n - 1) {
                return true;
            }
            if (step == 0) {
                return false;
            }
            int max = 0;
            int index = -1;
            for (int i = 1; i <= step; i++) {
                int value = i + nums[current + i];
                if (value >= max) {
                    max = value;
                    index = current + i;
                }
            }
            if (max == 0) {
                return false;
            }
            current = index;
        }
    }
}

この方法と動的計画法の違いといえば、 dp というキャッシュがないことで、利用するメモリの領域が固定だということです。単純にコードを見ると、最悪の場合でも計算量が O(N^2) になります。実際動くと動的計画法より早い感じですが、まだ一番理想的な方法ではありません。

問題の特徴を踏まえて、最初からではなく、最後から逆算してみます。

配列 [S, .... , Ex, Ey, Ez, E] に対して、最後の要素Eに辿り着くには、 Ex、Ey、Ezとそれ以前の要素からのが可能です。もう一つわかるのはEyからEzまでJUMPできるかつEzからEまでJUMPできると、EyからEまでJUMPできます。つまり

Ey -> Ez && Ez -> E => Ey -> E

逆だと成立しないです。ただ条件をつけば、成立します。

Ey -> E => Ey -> Ez && Ez -> E && Ez > 0

なぜなら、EyからEまでJUMPできると、もちろんEzまでもJUMPできる。ただEzからEまでJUMPできるのは、Ezが0以上の必要があります。

問題に戻ると、最初の要素から最後の要素までルートがあると、このルートは必ずEy、Ezを通るか。

考えてみてください。もしルートがEyを通ると、↑のルール条件つき（Ez > 0）でEzも通れるようになります。また、ルートがExを通ると、EyとEzの中で0以上の要素にも通れるようになります。

以上の現象に基づく解答は

class Solution {
    public boolean canJump(int[] nums){
        int n = nums.length;
        if(n == 0) {
            return false;
        }
        if(n == 1) {
            return true;
        }
        int tail = n - 1;
        for(int i = n - 2; i >= 0; i--) {
            if(nums[i] + i >= tail) {
                tail = i;
            }
        }
        return tail == 0;
    }
}

コードが簡単ですが、理解するには時間かかると思います。ここで tail を最後に通った要素の位置にします。 tail の初期値を最後の要素の位置にします。配列の最後から順に、各要素の位置にその値を足して tail に届くならその要素の位置で tail を更新していきます。最後に tail が 0 の場合は最初から最後までJUMPできるとなります。

いかがでしょうか。一つ簡単そうな問題でも、いろんな考えを元に解決できます。もし時間あれば、ぜひ試してください。