【Python】合計がpで割り切れるようにする最小の部分配列を求める方法

問題の概要

配列 nums と整数 p が与えられたとき、「残りの要素の合計が p で割り切れる」状態を作るために、最小の部分配列(ただし配列全体は削除しない)を取り除くことを考えます。求めたいのは削除すべき最小の部分配列の長さであり、そのような部分配列が存在しない場合は -1 を返します。

例えば、nums = [8,2,6,5,3]、p = 7 の場合を考えてみましょう。このとき出力は 1 になります。なぜなら、末尾の要素「3」を削除すれば残りの合計は 21 となり、21 は 7 で割り切れるからです。

解法のアイデア:累積和と剰余の活用

この問題は、累積和(プレフィックスサム)の剰余を利用することで効率的に解けます。ある区間 [i+1, j] を削除したとき、残りの合計が p で割り切れるという条件は、「j 番目までの累積和の剰余」と「i 番目までの累積和の剰余」の差が、配列全体の合計の剰余 s と一致すると言い換えられます。各剰余が最後に現れた位置を辞書に記録しておけば、配列を一度走査するだけで答えが求まります。

アルゴリズムの手順

• ans := 無限大(答えの初期値)
• s := nums の全要素の合計を p で割った余り
• d := {0: -1} というキーと値のペアを持つ辞書(マップ)
• cum := 0(累積和)
• s が 0 と等しい場合は、そのまま 0 を返す(既に条件を満たしているため)
• i を 0 から nums のサイズ未満まで繰り返す:
  • cum := cum + nums[i]
  • r := cum を p で割った余り
  • (r - s) mod p が d に存在する場合:
    • ans := ans と (i - d[(r-s) mod p]) の小さい方
  • d[r] := i
• ans が nums のサイズより小さければ ans を、そうでなければ -1 を返す

Pythonでの実装例

それでは、実際のコードを見て理解を深めましょう。

def solve(nums, p):
    ans = float("inf")
    s = sum(nums) % p
    d = {0:-1}
    cum = 0
    if s == 0:
        return 0
    for i in range(len(nums)):
        cum += nums[i]
        r = cum % p
        if (r-s) % p in d:
            ans = min(ans, i-d[(r-s)%p])
        d[r] = i
    return ans if ans<len(nums) else -1

nums = [8,2,6,5,3]
p = 7
print(solve(nums, p))

入力

[8,2,6,5,3], 7

出力

1

コードのポイント

このアルゴリズムの計算量は O(n)、必要な追加メモリは O(min(n, p)) であり、全ての部分配列を調べる総当たり法(O(n²))よりも大幅に効率的です。辞書 d には各剰余が最後に現れたインデックスを保存するため、常に「最も短くなる削除区間」を正しく検出できます。また、初期値 {0: -1} を入れておくことで、配列の先頭から削除するケースも漏れなく処理できる点が重要です。