現在の合計値でリストの要素を更新してターゲット配列に到達できるか判定するPythonプログラム

数値のリスト target が与えられたとします。ここで、与えられたリストと同じ長さを持ち、すべての要素が 1 で埋められたリスト X を考えます。次の操作を何度でも実行できます。

操作： X の任意のインデックス i を選び、X[i] を X の現在の合計値で置き換える。

この操作を繰り返した結果、X を target に変換できるかどうかを判定するのが本記事の目的です。

例えば、入力が target = [5, 9, 3] の場合、出力は True になります。その理由を見てみましょう。

• 初期状態：X = [1, 1, 1]
• 合計値 3 で更新 → [1, 1, 3]
• 合計値 5 で更新 → [5, 1, 3]
• 合計値 9 で更新 → [5, 9, 3]（ターゲットに一致！）

解決アプローチ

この問題を解くために、以下の手順に従います。

• nums の要素が 1 つだけの場合、その要素が 1 であれば true を返します
• q := nums のすべての数値を負の値にしたキューを作成します
• q をヒープ化します
• s := nums の全要素の合計とします
• ok := True とします
• ok が True である間、以下を繰り返します
  • x := ヒープから最小の要素（負の値のため実質最大値）を取り出し、符号を反転します
  • d := s - x（最大値以外の要素の合計）とします
  • d > 1 の場合は x2 := x mod d、それ以外の場合は 1 とします
  • s := s + x2 - x として合計を更新します
  • ok := x と x2 が異なるかどうかとします
  • x := x2 とします
  • -x をヒープ q に挿入します
• 最後に、q のすべての要素が -1 であれば true を返します

アルゴリズムのポイント

この手法の鍵は「逆算」の考え方です。ターゲット配列の最大値は、必ず直前の操作で「それ以外の要素の合計」から生成されたものです。そこで、最大値 x から残りの合計 d を引いて余りを求めることで、一歩前の状態へ効率よく戻れます。さらに mod 演算を使うことで、巨大な数値が含まれる場合でも複数ステップを一度に処理でき、高速に判定できるのが特徴です。

実装例

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

class Solution:
   def solve(self, nums):
      if len(nums) == 1:
         return nums == [1]
      from heapq import heapify, heappop, heappush

      q = [-x for x in nums]
      heapify(q)
      s = sum(nums)
      ok = True

      while ok:
         x = -heappop(q)
         d = s - x
         x2 = x % d if d > 1 else 1
         s += x2 - x
         ok = x != x2
         x = x2
         heappush(q, -x)

      return all(x == -1 for x in q)

ob = Solution()
target = [5, 9, 3]
print(ob.solve(target))

入力

[5, 9, 3]

出力

True