Pythonで公正な配列を作れる削除インデックスの数を求めるプログラム

問題の概要

配列 nums が与えられたとします。私たちはちょうど1つのインデックスを選び、その位置にある要素を削除できます(削除後は残りの要素が前に詰まり、インデックスが変化することに注意してください)。

ここで、偶数番目(インデックス0, 2, 4, …)の値の合計と奇数番目(インデックス1, 3, 5, …)の値の合計が等しいとき、その配列を「公正(fair)」であると呼びます。求めたいのは、1つの要素を削除した結果として配列が公正になるような、インデックスの選び方の総数です。

入力例と考え方

たとえば入力が nums = [5,3,7,2] の場合、出力は 1 になります。各インデックスを削除した場合を順に確認してみましょう。

• インデックス0を削除 → 配列は [3,7,2]。偶数番目の合計: 3+2 = 5、奇数番目の合計: 7(不公正)

• インデックス1を削除 → 配列は [5,7,2]。偶数番目の合計: 5+2 = 7、奇数番目の合計: 7(公正)

• インデックス2を削除 → 配列は [5,3,2]。偶数番目の合計: 5+2 = 7、奇数番目の合計: 3(不公正)

• インデックス3を削除 → 配列は [5,3,7]。偶数番目の合計: 5+7 = 12、奇数番目の合計: 3(不公正)

公正になるのは「インデックス1を削除した場合」だけなので、答えは 1 となります。

アルゴリズム

毎回新しい配列を作り直して合計を計算する素朴な方法では O(n²) の時間がかかりますが、削除による偶奇の入れ替わりに注目すれば、O(n) で効率よく解けます。手順は以下の通りです。

• res := 0、sm1 := 0、sm2 := 0 で初期化する
• i を 1 から nums のサイズ − 1 まで動かしながら、
  • i が偶数なら sm1 := sm1 + nums[i]
  • そうでなければ sm2 := sm2 + nums[i]
• sm1 と sm2 が等しければ res := res + 1(これは「インデックス0を削除した場合」に相当)
• 再び i を 1 から nums のサイズ − 1 まで動かしながら、
  • i が奇数なら sm1 := sm1 − nums[i] + nums[i−1]
  • そうでなければ sm2 := sm2 − nums[i] + nums[i−1]

  そして sm1 と sm2 が等しければ res := res + 1

• 最後に res を返す

なぜこれでうまくいくのか

最初のループでは、「インデックス0を削除した場合」の偶数番目・奇数番目それぞれの合計を求めています。先頭を削除すると、それ以降の要素はすべて1つ前に詰まるため偶奇が反転します。つまり、元の配列で奇数インデックスだった要素(sm1)が新しい偶数番目へ、偶数インデックス(2以降)だった要素(sm2)が新しい奇数番目へ移ります。

2つ目のループでは、削除位置を1つずつ右へずらしていきます。削除位置を i−1 から i に移すと、それまで配列から外れていた nums[i−1] が復活し、代わりに nums[i] が取り除かれます。そこで、nums[i] の偶奇に応じて該当する合計から nums[i] を引き、nums[i−1] を加えれば、各ステップの更新は O(1) で済みます。

実装例(Python)

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

def solve(nums):
    res, sm1, sm2 = 0, 0, 0
    for i in range(1, len(nums)):
        if i%2 == 1:
            sm1 += nums[i]
        else:
            sm2 += nums[i]
    if sm1 == sm2:
        res += 1

    for i in range(1, len(nums)):
        if i % 2 == 1:
            sm1 = sm1 - nums[i] + nums[i-1]
        else:
            sm2 = sm2 - nums[i] + nums[i-1]
        if sm1 == sm2:
            res += 1

    return res

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

入力

[5,3,7,2]

出力

1

計算量

配列の走査は全体で2回だけなので、時間計算量は O(n) です。また、使用する変数は定数個のみのため、空間計算量も O(1) で抑えられます。