部分配列の和が0になる個数をO(N)で求める

広告

面接のクイズで次のような問題が出ました。

配列の部分配列の和が0になる個数を求めよ。
部分配列とは、例えば
int arr[] =  {2, -2, 3, 0, 4, -7};
の[2..5]は{3,0,4,-7}で合計すると0になる。

この例では

[0..1] = {2, -2}
[3..3] = {0}
[0..5] = {2, -2, 3, 0, 4, -7} // 全部
[2..5] = {3, 0, 4, -7}

が求める配列で、個数は4つ。

計算量はO(N log N)で求めよ。

O(N^2)の解答例

ものすごく簡単に書くと

int count(vector<int> &A)
{
    int count = 0;

    for( unsigned int i = 0; i < A.size(); i++ )
    {
        int sum = 0;

        for( unsigned int j = i; j < A.size(); j++ )
        {
            sum += A[j];
            if( sum == 0 )
            {
                count++;
                printf("[%d..%d]\n", i, j);
            }
        }
    }

    return count;
}

実行例

[0..1]
[0..5]
[2..5]
[3..3]
4

答えはあっていますが、O(N^2)なので大きなケースでは落ちます。

O(N)の解答例

頭から現在地点までの和を計算します。

index 0 1 2 3 4 5
2 -2 3 0 4 7
0..iまでの和 2 0 3 3 7 0

これを見ると次のことがわかります。

0..iまでの和が0の場合

i = 1とi = 5で0になっています。これは

  • A[0] + A[1] = 0
  • A[0] + A[1] + A[2] + A[3] + A[4] + A[5] = 0

ということを表していますので、単純にカウントに加えてOKです。

同じ数字が複数出る場合

  • 和が3: i = 2, i = 3
  • 和が0: i = 1, i = 5

が同じ値になっています。これはつまり、最初に出現したときから、値を足していって結局同じ値になったことを意味します。

i = 1で0が出現し、i = 5でも同じ値の0が出現したということは、[0, 1]の和 = [0, 1, 2, 3, 4, 5]の和なので、[2, 3, 4, 5]の和も0ということになります。これもカウントする必要があります。

これをコードで書くと

int count2(vector<int> &A)
{
    map<int, set<int>> s;
    int count = 0;

    // Traverse throught array and store prefix sums
    int sum = 0;
    for (int i = 0 ; i < A.size() ; i++)
    {
        sum += A[i];

        if (sum == 0)
        {
            count++;
            printf("[%d..%d]\n", 0, i);
        }
        if( s.count(sum) != 0 )
        {
            for( int j: s[sum] )
            {
                count++;
                printf("[%d..%d]\n", j + 1, i);
            }
        }
        s[sum].insert(i);
    }
    return count;
}

例えばこのような感じになります。連想配列mapに

sum -> これまでに出現した位置

を記録しています。ただ数を数えればいいだけなら過去の出現数を入れておけばいいでしょう。これはループが1重なので計算量はO(N)で、要求された計算量より小さくなりました。

O(N log N)を求められたので二分検索のようなものを考えてしまいましたが、O(N)のほうが簡単に思いつきそうです。

応用問題

和が0の場合でしたが、和がxの場合はもうひと工夫いるでしょう。

部分配列の和が最大になる箇所を求める場合ももう少し複雑になります。

参考サイト

http://stackoverflow.com/questions/5534063/zero-sum-subarray