Pythonで合計がkとなる重複しない2つの部分リストの長さの合計を求めるプログラム

数値のリスト nums と別の値 k が与えられます。nums の中から互いに重ならない2つの部分リスト（サブリスト）を見つけ、それぞれの要素の合計が k と等しくなるようにします。そして、その2つの部分リストの長さの合計を求めるのがこの問題です。

条件を満たす組み合わせが複数存在する場合は、最も短い2つの部分リストを選び、その長さの合計を返します。条件を満たす組み合わせが見つからない場合は -1 を返します。

具体例

たとえば、入力が次の場合を考えてみましょう。

nums = [7, 10, -2, -1, 4, 3]、k = 7

このとき出力は 3 になります。[7]（長さ1）と [4, 3]（長さ2）という2つの部分リストを選べば、どちらも合計が 7 になり、長さの合計は 1 + 2 = 3 となるためです。同じく合計が 7 になる [10, -2, -1] は長さが 3 と長いため、より短い組み合わせがあるこちらは採用しません。

解法のアプローチ

この問題は、左端からの最短部分リスト長（prefix）と右端からの最短部分リスト長（suffix）を前計算しておき、分割位置ごとに組み合わせて最小値を求めることで効率的に解けます。手順は以下の通りです。

• N := 配列 A のサイズとする

• prefix := サイズ N の配列を無限大で初期化する

• last := マップ {0: -1} とする

• s := 0 とする

• i を 0 から N-1 までループ：

  • s := s + A[i]

  • prefix[i] := i - last[s - target]（キーが存在しない場合は -無限大）

  • last[s] := i

• i を 1 から N-1 までループ：

  • prefix[i] := min(prefix[i], prefix[i - 1])

• suffix := サイズ N の配列を無限大で初期化する

• last := マップ {0: N} とする

• s := 0 とする

• i を N-1 から 0 まで減らしながらループ：

  • s := s + A[i]

  • suffix[i] := last[s - target] - i（キーが存在しない場合は 無限大）

  • last[s] := i

• i を N-2 から 0 まで減らしながらループ：

  • suffix[i] := min(suffix[i], suffix[i + 1])

• ans := 各 i（0 ≤ i ≤ N-2）について prefix[i] + suffix[i + 1] の最小値

• ans が無限大より小さければ ans を、そうでなければ -1 を返す

実装例

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

class Solution:
    def solve(self, A, target):
        INF = float("inf")
        N = len(A)
        prefix = [INF] * N
        last = {0: -1}
        s = 0
        for i in range(N):
            s += A[i]
            prefix[i] = i - last.get(s - target, -INF)
            last[s] = i
        for i in range(1, N):
            prefix[i] = min(prefix[i], prefix[i - 1])
        suffix = [INF] * N
        last = {0: N}
        s = 0
        for i in range(N - 1, -1, -1):
            s += A[i]
            suffix[i] = last.get(s - target, INF) - i
            last[s] = i
        for i in range(N - 2, -1, -1):
            suffix[i] = min(suffix[i], suffix[i + 1])
        ans = min(prefix[i] + suffix[i + 1] for i in range(N - 1))
        return ans if ans < INF else -1
ob = Solution()
nums = [7, 10, -2, -1, 4, 3]
k = 7
print(ob.solve(nums, k))

入力

[7, 10, -2, -1, 4, 3], 7

出力

3

計算量について

各ループは配列を高々一度ずつ走査するだけなので、時間計算量は O(N) です。また、補助配列 prefix・suffix とハッシュマップを使用するため、空間計算量も O(N) となります。累積和とハッシュマップを組み合わせることで、「合計が target に一致する部分リスト」の探索を高速に行えるのがこの手法のポイントです。