Pythonで2つの配列の合計を一致させる最小操作回数を求めるプログラム

問題の概要

nums1 と nums2 という2つの整数配列が与えられます。各配列の要素はすべて 1〜6 の範囲の値です。1回の操作では、どちらかの配列内の任意の要素を、1〜6 の範囲内の任意の値に書き換えることができます。

このとき、nums1 の要素の合計と nums2 の要素の合計を一致させるのに必要な最小の操作回数を求めます。どうしても一致させられない場合は -1 を返してください。

例えば、入力が nums1 = [1, 5, 6]、nums2 = [4, 1, 1] の場合、答えは 2 になります。実際、次のように2回の操作で合計を揃えられます。

• 1回目:nums2 を [4, 1, 1] → [4, 1, 6] に変更(合計 6 → 11)
• 2回目:[4, 1, 6] → [4, 2, 6] に変更(合計 11 → 12)

これで nums2 の合計は nums1 の合計(12)と一致します。

解法のアプローチ:貪欲法

この問題は貪欲法(グリーディ法)で効率的に解くことができます。基本戦略は次のとおりです。

• 合計が小さい方の配列は、小さい値から順に「6」へ引き上げて合計を増やす。
• 合計が大きい方の配列は、大きい値から順に「1」へ引き下げて合計を減らす。
• 各ステップでは、「6への引き上げによる増分」と「1への引き下げによる減分」のうち、効果の大きい方を採用する。

アルゴリズムの手順

1. 両配列の合計 s1、s2 を計算する。
2. nums1 と nums2 をそれぞれ昇順にソートする。
3. s1 > s2 の場合は nums1 と nums2(および s1 と s2)を入れ替え、以降は常に s1 ≤ s2 が成り立つようにする。
4. ポインタ left = 0、right = len(nums2) − 1 を初期化する。
5. s1 == s2 になった時点で操作回数 ans を返す。まだ一致していなければ、次の判定を行う。
   • curr_left = nums1[left](left が範囲外なら 7)
   • curr_right = nums2[right](right が範囲外なら 0)
   • 6 − curr_left ≥ curr_right − 1 ならば:s1 に min(6 − curr_left, s2 − s1) を加算し、left を +1 する。
   • それ以外の場合:s2 から min(curr_right − 1, s2 − s1) を減算し、right を −1 する。
6. 1ステップごとに ans を +1 し、ループを抜けた時点で s1 != s2 なら -1 を返す。

ここで、範囲外を表す番兵値として curr_left = 7(増分 6 − 7 = −1)や curr_right = 0(減分 0 − 1 = −1)を使うことで、使い切った側の候補が自動的に選ばれないようになっているのがポイントです。

Pythonでの実装例

def solve(nums1, nums2):
    s1 = sum(nums1)
    s2 = sum(nums2)
    nums1.sort()
    nums2.sort()
    if s1 > s2:
        nums1, nums2 = nums2, nums1
        s1, s2 = s2, s1

    ans = 0
    left, right = 0, len(nums2) - 1
    while left < len(nums1) or right >= 0:
        if s1 == s2:
            return ans
        curr_left = nums1[left] if left < len(nums1) else 7
        curr_right = nums2[right] if right >= 0 else 0
        if 6 - curr_left >= curr_right - 1:
            s1 += min(6 - curr_left, s2 - s1)
            left += 1
        else:
            s2 -= min(curr_right - 1, s2 - s1)
            right -= 1
        ans += 1
    return -1 if s1 != s2 else ans

nums1 = [1, 5, 6]
nums2 = [4, 1, 1]
print(solve(nums1, nums2))

入力

[1, 5, 6], [4, 1, 1]

出力

2

計算量について

ソートに O(n log n)、その後の貪欲ループは各要素を最大1回ずつ処理するだけなので O(n)。全体の時間計算量は O(n log n)(n は配列の長さ)です。

なお、片方の配列が極端に短い場合(例：長さ1の配列の最大合計 6 より、長さ10の配列の最小合計 10 が大きい場合など)には、すべての要素を使い切っても合計を一致させることができません。このようなケースでも、本アルゴリズムは正しく -1 を返します。