Pythonで2つの配列から合計が等しい部分配列を検索する方法

問題概要

大きさNの2つの配列PとQがあり、それぞれ1からNまでの数値を含んでいると仮定します。この問題では、与えられた2つの配列から合計が等しくなる部分配列（サブ配列）を見つけ出し、そのインデックスを返す必要があります。解が存在しない場合は-1を返します。

例えば、入力が以下の場合を考えてみましょう。

• P = [2, 3, 4, 5, 6]
• Q = [9, 3, 2, 6, 5]

この場合の出力は「最初の配列のインデックス：0, 1, 2」「2番目の配列のインデックス：0」となります。これは、P[0..2] = 2 + 3 + 4 = 9 と Q[0] = 9 の合計が一致しているためです。

アルゴリズムの手順

この問題は、累積和とハッシュマップ（辞書）を組み合わせることで効率的に解けます。具体的な手順は以下の通りです。

1. get_subarray() 関数を定義します。引数としてP、Q、swapを受け取ります。
2. N := Pのサイズ
3. index := 新しいマップ（辞書）を作成
4. difference := 0、j := 0 で初期化
5. index[0] := (-1, -1) というペアを設定
6. i を 0 から N まで繰り返します。
   • Q[j] < P[i] の間、j を増加させ続ける
   • difference := Q[j] - P[i]
   • difference が index に存在する場合：
     • swap が True の場合：
       • idx := index[Q[j] - P[i]]
       • P については idx[1]+1 から j までの値を表示
       • Q については idx[0]+1 から i までの値を表示
     • それ以外の場合：
       • idx := index[Q[j] - P[i]]
       • P については idx[0]+1 から i までの値を表示
       • Q については idx[1]+1 から j までの値を表示
     • return で終了
   • index[difference] := (i, j) を記録
7. 見つからなかった場合は -1 を表示

メイン処理の流れ

メインメソッドでは、以下の処理を行います。

• P と Q をそれぞれ累積和に変換して更新する
• N := Pのサイズを取得
• Q[N-1] > P[N-1] の場合：get_subarray(P, Q, False) を呼び出す
• それ以外の場合：get_subarray(Q, P, True) を呼び出す

実装例

理解を深めるために、以下のPythonコードを実装してみましょう。

def show_res(x, y, num):
    print("Indices of array", num, ":", end=" ")
    for i in range(x, y):
        print(i, end=", ")
    print(y)

def get_subarray(P, Q, swap):
    N = len(P)
    index = {}
    difference, j = 0, 0
    index[0] = (-1, -1)
    for i in range(0, N):
        while Q[j] < P[i]:
            j += 1
        difference = Q[j] - P[i]
        if difference in index:
            if swap:
                idx = index[Q[j] - P[i]]
                show_res(idx[1] + 1, j, 1)
                show_res(idx[0] + 1, i, 2)
            else:
                idx = index[Q[j] - P[i]]
                show_res(idx[0] + 1, i, 1)
                show_res(idx[1] + 1, j, 2)
            return
        index[difference] = (i, j)
    print(-1)

def cumsum(arr):
    n = len(arr)
    for i in range(1, n):
        arr[i] += arr[i - 1]

P = [2, 3, 4, 5, 6]
Q = [9, 3, 2, 6, 5]
cumsum(P)
cumsum(Q)
N = len(P)
if Q[N - 1] > P[N - 1]:
    get_subarray(P, Q, False)
else:
    get_subarray(Q, P, True)

入力

[2, 3, 4, 5, 6],[9, 3, 2, 6, 5]

出力

Indices of array 1 : 0, 1, 2
Indices of array 2 : 0

アルゴリズムのポイント

このアルゴリズムの鍵となるのは累積和の活用です。各配列を累積和に変換することで、「任意の区間の合計」を「2つの累積和の差」として表現できます。そのため、2つの配列の累積和の差分が同じになる位置を記録しておき、同じ差分が再び現れた時点で、合計が等しい部分配列のペアを即座に特定できます。

また、swapフラグを使って引数の順序を制御することで、どちらの配列の合計が大きいかに関係なく、同じ関数で処理できる点も実装上の工夫です。計算量は線形時間O(N)で抑えられるため、大きなデータセットに対しても高速に動作します。