🤏

[leetCode] Top Interview 150: Product of Array Except Self

2024/01/10に公開

Java

LeetCode

tech

リンク

概要

数列numsが与えられる
nums[i]に対して、それ以外のすべての要素を掛け合わせた積をそれぞれ求める
- たとえば[1, 2, 3, 4]だと[24, 12, 8, 6] みたいになる
$O(N)$ で実装する必要がある
除算を使用してはいけない

解法

さきに、めっちゃ参考になるノート（英語）を共有しておきます。
たぶんこっちを読んだほうが早いと思います。

総当たり

さて、問題について考えていきます。
まず、この問題の意図を明らかにするため、もっとも直感的な方法で実装してみます。

class Solution {
    public int[] productExceptSelf(int[] nums) {
        int n = nums.length;
        int ans[] = new int[n];
        
        for(int i = 0; i < n; i++) {
            int prod = 1;
            for(int j = 0; j < n; j++) {
                if(i == j) continue;   # 選択した要素は掛け合わせない
                pro *= nums[j];   # それ以外は掛け合わせる
            }
            ans[i] = prod;
        }
        
        return ans;
    }
}

最もシンプルな、総当たりでの求め方です。
ひとつめのfor文で各要素を選択し、その中のfor文で選択した要素以外の要素をすべて掛け合わせることで求めているのがわかります。

ただし、このままでは計算量が $O(N^2)$ になってしまいますね。
この問題では、より計算量を意識した実装が必要になりそうです。

すべての要素の積から除算する

そこで、ちょっと違った視点から考えてみます。
「nums[i]以外の要素の積」は、「すべての要素の積　÷　nums[i]」であることがわかるでしょうか。

例えば、[1, 2, 3, 4]という数列について考えてみます。

i=1 のとき、nums[i] は 2 である
nums[i]以外の積は 1 x 3 x 4 = 12 となる
これは、numsのすべての要素の積 = 24 から nums[i] = 2 を割った数に等しい (24 ÷ 2 = 12）

というように、最初にすべての要素をかけ合わせておいて、特定の要素を抜きたければ、その要素だけを割ればよいのです。
というのを実装したのが以下のコード。

（このコメントを参考にさせていただきました。ありがとうございます）

class Solution {
    public int[] productExceptSelf(int[] nums) {
        int n = nums.length;
        int allProdWithoutZero = 1;
        int countZeroes = 0;
        int zeroIndex = -1;
        
        // 0以外のすべての要素を掛け合わせる
        for(int i=0; i<n; i++) {
            int num = nums[i];

            // 0の場合はカウントして次へ
            if(num == 0) {
                countZeroes++;
                zeroIndex = i;
                continue;
            }
            allProdWithoutZero *= num;
        }

        // int配列には初期値として 0 が入る
        int[] answer = new int[n];

        // 0が存在しない場合: 各要素を除算して答えを求める
        if(countZeroes == 0) {
            for(int i=0; i<n; i++) {
                answer[i] = allProdWithoutZero / nums[i];
            }

        // 0が1つだけある場合: 0以外の要素の答えは0になる
        } else if(countZeroes == 1) {
            answer[zeroIndex] = allProdWithoutZero;
        }

        // 0が2つ以上ある場合: すべての答えが0になる

        return answer;
    }
}

ただし、今回の問題の制限として、

除算を使用してはいけない

というものがあるので、これも解決策としては外れます。

配列を「要素の前」と「要素の後」に分割して求める

さて、本題の方法です。

最初の方法（総当たり）では、各要素を選択し、選択した要素以外の要素をすべて掛け合わせることで要素を求めました。
この方法では $O(N^2)$ がかかってしまい、条件を満たせませんでしたね。

ここで、少し考え方を変えてみます。
先ほどと同じように、[1, 2, 3, 4]という配列について考えてみましょう。

i=1 のとき、 nums[i] は 2 である
nums[i]以外の積は 1 x 3 x 4 = 12 となる
この計算は、以下のように分割できる
a. i < 1 の要素のすべての積: 1
b. i > 1 の要素のすべての積: 3 x 4 = 12
c. これらを掛け合わせると 1 x 12 = 12 となり、答えに等しい

この時点では「なぜ分割して求めるの？」となりますが、実は以上の考え方を活用することで、 $O(N)$ で計算ができるようになります。
以下のコードを見てください。

class Solution {
    public int[] productExceptSelf(int[] nums) {
        int n = nums.length;
        int[] pre = new int[n];
        int[] suf = new int[n];

        // pre[i] は、「nums[i] より前のすべての要素の積」を表す
        pre[0] = 1;
        for(int i=1; i<n; i++) {
            pre[i] = pre[i-1] * nums[i-1];
        }

        // suf[i] は、「nums[i] より後のすべての要素の積」を表す
        suf[n-1] = 1;
        for(int i=n-2; i>=0; i--) {
            suf[i] = suf[i+1] * nums[i+1];
        }

        // pre[i] * suf[i] = answer[i] になる
        int[] answer = new int[n];
        for(int i=0; i<n; i++) {
            answer[i] = pre[i] * suf[i];
        }

        return answer;
    }
}

このコードでは、以下のような感じの計算をしています。

pre("prefix" の意です) と suf("suffix" の意です) という2つの配列を用意
- pre[i] は、 「nums[i] より前のすべての要素の積」 を示す
- suf[i] は、 「nums[i] より後のすべての要素の積」 を示す

ここで、さっきの説明を思い出してください。

i=1 のとき、 nums[i] は 2 である

nums[i]以外の積は 1 x 3 x 4 = 12 となる

この計算は、以下のように分割できる
a. i < 1 の要素のすべての積: 1
b. i > 1 の要素のすべての積: 3 x 4 = 12
c. これらを掛け合わせると 1 x 12 = 12 となり、答えに等しい

やっていることはこれと全く同じ。
ただ、preとsufという配列に、積み上げるように計算結果をメモし、あとで答えを効率的に計算しているだけです。

もちろん二重ループなどは使っていないので、時間計算量は $O(N)$ となります。

（おまけ）↑を空間計算量 $O(1)$ で実装する

問題文の端に、こんなことが書いてあります。

Follow up: Can you solve the problem in O(1) extra space complexity? (The output array does not count as extra space for space complexity analysis.)

要するに、「この問題は空間計算量が $O(1)$ で実装できるで（答えを記録する配列は特別に許すで）」という意味です。
「空間計算量が $O(1)$ 」というのは、言い換えれば 「preやsufなどの配列を使わないように実装し直す」 ということですね。

そこで、preやsufを求める過程を配列に記録するのではなく、直接answerにかけ合わせる形で実装してみます。
実装したものがこちら。

class Solution {
    public int[] productExceptSelf(int[] nums) {
        int n = nums.length;
        int[] answer = new int[n];
        Arrays.fill(answer, 1);   // 直接掛け合わせるため、すべての要素を 1 で初期化しておく
        int current = 1;

        // pre
        for(int i=0; i<n; i++) {
            answer[i] *= current;
            current *= nums[i];
        }

        // suf
        current = 1;
        for(int i=n-1; i>=0; i--) {
            answer[i] *= current;
            current *= nums[i];
        }

        return answer;
    }
}

だいぶすっきりしましたね。
時間計算量は変わらず $O(N)$ 、空間計算量は $O(1)$ となりました。

おわりに

日課にしては規格外のボリュームで書いてしまい、ほかの予定に費やす時間がお釈迦になりました。まあいいか。