Pythonで配列を等しい合計のサブ配列に分割できる合計値を見つける方法

整数の配列Aが与えられたとき、ある値sum[i]ごとに、配列を合計がsum[i]となる複数のサブ配列に分割できるような、すべての合計値を見つける必要があります。もし配列を等しい合計のサブ配列に分割できない場合は、-1を返します。

例えば、入力が A = [2, 4, 2, 2, 2, 4, 2, 6] の場合、出力は [6, 8, 12] になります。これは、配列を合計が6、8、12となるサブ配列にそれぞれ分割できるためです。具体的な分割例は以下の通りです。

• 合計6の場合： {2, 4}, {2, 2, 2}, {4, 2}, {6}
• 合計8の場合： {2, 4, 2}, {2, 2, 4}, {2, 6}
• 合計12の場合： {2, 4, 2, 2, 2}, {4, 2, 6}

解決のためのアプローチ

この問題は、以下の手順で解くことができます。

1. 累積和テーブルの作成： 配列aのサイズをnとし、サイズnのテーブルを作成します。table[0]にa[0]を設定し、以降は「table[i] = a[i] + table[i-1]」という式で累積和を順番に計算していきます。
2. 全体の合計Sの取得： テーブルの最後の要素table[n-1]が、配列全体の合計Sとなります。
3. マップへの記録： 累積和のすべての値をマップ(my_map)に登録します。これにより、任意の値が累積和として存在するかどうかをO(1)で高速に確認できるようになります。
4. 約数の探索： 1から√Sまでの各数値iについて、Sがiで割り切れるかどうかを確認します。割り切れる場合は、iとS÷iの両方がSの約数となります。
5. 分割可能性の検証： 各約数part_1(およびpart_2)について、「part_1、2×part_1、3×part_1、…」とSまでのすべての倍数がマップに存在するかを確認します。すべて存在すれば、その値を合計とするサブ配列への分割が可能です。
6. 結果の返却： 有効な合計値がひとつも見つからない場合は-1を返し、見つかった場合はそれらの集合を返します。

実装例

以下のPythonコードを実装例として見てみましょう。

from math import sqrt

def find_sum(a):
    n = len(a)
    table = [0] * n
    table[0] = a[0]
    for i in range(1, n):
        table[i] = a[i] + table[i - 1]
    S = table[n - 1]
    my_map = {}
    for i in range(n):
        my_map[table[i]] = 1
    answer = set()
    for i in range(1, int(sqrt(S)) + 1):
        if (S % i == 0):
            is_present = True
            part_1 = i
            part_2 = S // i
            for j in range(part_1, S + 1, part_1):
                if j not in my_map:
                    is_present = False
                    break
            if (is_present and part_1 != S):
                answer.add(part_1)
            is_present = True
            for j in range(S // i, S + 1, S // i):
                if j not in my_map:
                    is_present = False
                    break
            if (is_present and part_2 != S):
                answer.add(part_2)
    if (len(answer) == 0):
        return -1
    return answer

a = [2, 4, 2, 2, 2, 4, 2, 6]
print(find_sum(a))

入力

[2, 4, 2, 2, 2, 4, 2, 6]

出力

{8, 12, 6}

アルゴリズムのポイント

このアルゴリズムを理解する上で重要な考え方は以下の通りです。

• 配列を合計kのサブ配列に分割できるのは、k、2k、3k、…という値がすべて累積和として現れる場合に限られます。これは、各サブ配列の区切り位置が必ずこれらの累積和の位置と一致する必要があるためです。
• 候補をSの約数だけに絞ることで、探索範囲を大幅に削減できます。約数はペア(i, S/i)で現れるため、√Sまで調べるだけで十分です。
• part_1やpart_2がS自身と等しいケースは除外します。これは配列全体を1つのサブ配列とみなすことを意味し、実際の「分割」にはならないためです。