Pythonでクエリ処理後の偶数の合計を効率的に求める方法

整数の配列 A と、クエリを格納した配列 queries があるとします。i番目のクエリでは、value = queries[i][0]、index = queries[i][1] となり、A[index] に value を加算します。そして、i番目のクエリに対する答えは、更新後の配列 A に含まれる偶数の合計値です。すべてのクエリに対する答えを順番に求め、それらを配列として返すのがこの問題の目的です。

問題の例

例として、配列が [1,2,3,4]、クエリ配列が [[1,0],[-3,1],[-4,0],[2,3]] の場合を考えてみましょう。このとき、答えの配列は [8,6,2,4] になります。

• 初期状態： 配列は [1,2,3,4]
• 1番目のクエリ： A[0] に 1 を加算すると、配列は [2,2,3,4] になります。偶数の合計は 2 + 2 + 4 = 8
• 2番目のクエリ： A[1] に -3 を加算すると、配列は [2,-1,3,4] になります。偶数の合計は 2 + 4 = 6
• 同様の手順を繰り返すことで、最終的に [8,6,2,4] という結果が得られます。

解法のアプローチ

各クエリのたびに配列全体を走査して偶数の合計を再計算すると、計算量が O(N×Q) になり非効率です。そこで、現在の偶数の合計を変数として保持し、更新された要素だけを差分として調整する方法が有効です。手順は以下の通りです。

• 結果を格納するための配列 res を定義する
• sum := 0 と初期化する
• 配列 A の各要素 i について：
  • i が偶数であれば、sum := sum + i とする
• queries の各クエリ i について：
  • index := i[1]
  • val := i[0]
  • A[index] が偶数であれば、まず sum := sum - A[index] として合計から除外する
  • A[index] := A[index] + val で値を更新する
  • 更新後の A[index] が偶数であれば、sum := sum + A[index] として合計に加える
  • sum を res に追加する
• res を返す

この手法により、各クエリの処理は O(1) で完了し、全体の計算量は O(N + Q) に抑えられます。

Pythonでの実装例

以下の実装を見ると、理解がより深まるでしょう。

class Solution(object):
   def sumEvenAfterQueries(self, A, queries):
      result = []
      sum = 0
      for i in A:
         if i%2==0:
            sum+=i
      for i in queries:
         index = i[1]
         val = i[0]
         if A[index] % 2==0:
            sum-=A[index]
         A[index]+=val
         if A[index]%2==0:
            sum+=A[index]
         result.append(sum)
      return result
ob1 = Solution()
print(ob1.sumEvenAfterQueries([1,2,3,4], [[1,0],[-3,1],[-4,0],[2,3]]))

入力

[1,2,3,4]
[[1,0],[-3,1],[-4,0],[2,3]]

出力

[8,6,2,4]