Pythonで配列から1つの要素を削除して「良い配列」になるインデックスをすべて見つける方法

問題の概要

数値の配列 A が与えられたとき、i 番目の要素を削除した後に「良い配列(good array)」となるような、すべてのインデックスを見つける必要があります。ここでの条件は以下の通りです。

• 良い配列とは、配列内のある要素が、それ以外のすべての要素の合計と等しい配列のことです。
• インデックスは 1 始まり(1-based)で表します。

たとえば、入力が [10, 4, 6, 2] の場合、出力は [1, 4] になります。

• A[1](=10)を削除すると、配列は [4, 6, 2] となり、6 = 4 + 2 が成立するため良い配列です。
• A[4](=2)を削除すると、配列は [10, 4, 6] となり、10 = 4 + 6 が成立するため良い配列です。

解法のアプローチ

この問題を効率的に解く鍵は、次の性質にあります。要素を1つ削除した後の残りの合計を S とすると、「ある要素 = 残りの要素の合計」という条件より、S は必ずその特別な要素の2倍になります。つまり、S が偶数であり、かつ S ÷ 2 の値が残りの配列の中に存在していればよいことになります。

具体的な手順は以下の通りです。

• n := 配列 A のサイズ
• add := 全要素の合計(初期値 0)
• my_map := 各要素の出現回数を記録するマップ(辞書)
• i を 0 から n−1 まで繰り返す:
  • my_map[A[i]] のカウントを 1 増やす
  • add に A[i] を加算する
• 再び i を 0 から n−1 まで繰り返す:
  • k := add − A[i](A[i] を除外した残りの合計)
  • k が偶数の場合:
    • k := k / 2
    • k が my_map に存在し、かつ(A[i] == k で my_map[k] > 1、または A[i] != k)を満たせば、i + 1 を出力する

「A[i] == k の場合は my_map[k] > 1 が必要」という条件があるのは、削除対象の要素自身を誤って答えとしてカウントしないためです。A[i] と k が同じ値でも、その値が配列内に2個以上存在すれば、削除後も別の要素として k が残るため、そのインデックスは有効になります。

実装例

以下に Python での実装を示します。

from collections import defaultdict
def find_indices(A):
    n = len(A)
    add = 0
    my_map = defaultdict(lambda:0)
    for i in range(n):
        my_map[A[i]] += 1
        add += A[i]
    for i in range(n):
        k = add - A[i]
        if k % 2 == 0:
            k = k >> 1
            if k in my_map:
                if ((A[i] == k and my_map[k] > 1) or (A[i] != k)):
                    print((i + 1))
A = [10, 4, 6, 2]
find_indices(A)

入力

[10, 4, 6, 2]

出力

1
4

計算量について

このアルゴリズムは、要素数を n とすると、前処理(合計と出現回数の集計)に O(n)、判定フェーズも各インデックスごとに定数時間の処理しか行わないため、全体で O(n) の計算量で動作します。ハッシュマップ(defaultdict)による高速な参照が、この効率性を実現するポイントです。