Pythonで隣接するペアの合計が完全平方数となる順列の数をカウントするプログラム

数値のリスト nums が与えられたとき、「隣接する任意の2つの値の合計が完全平方数(平方数)となる」ような順列の個数を求める問題を考えてみましょう。2つの順列 A と B は、あるインデックス i において A[i] と B[i] が異なるとき、互いに異なる順列として区別します。

例えば、入力が nums = [2, 9, 7] の場合、出力は 2 になります。これは [2, 7, 9](2+7=9、7+9=16)と [9, 7, 2](9+7=16、7+2=9)の2通りが条件を満たすためです。どちらも隣接要素の和が 9 や 16 といった完全平方数になっていますね。

解法のアプローチ

この問題はバックトラッキング(試行の巻き戻し)を使って解くのが効果的です。さらに、同じ値から始まる重複した探索を避けるために、各段階で「訪問済みの合計値」を記録する集合(visited set)による枝刈りを行います。これにより、無駄な再帰呼び出しを大幅に減らせます。

アルゴリズムの手順

• 結果を格納する変数 res を 0 で初期化します。
• 関数 util(i) を定義します。
• i + 1 が nums のサイズと等しい場合(末尾まで正しく配置できた場合):
  • res を 1 増やして処理を終了します。
• 空の集合 visited を作成します。
• j を i + 1 から nums のサイズまで繰り返します:
  • s = nums[i] + nums[j] を計算します。
  • s が未訪問であり、かつ s の平方根の2乗が s に一致する(s が完全平方数である)場合:
    • s を visited に追加します。
    • nums[i + 1] と nums[j] を入れ替えます。
    • util(i + 1) を再帰的に呼び出します。
    • nums[i + 1] と nums[j] を元に戻します(バックトラック)。
• メイン処理では以下を実行します:
• 新しい集合 visited を作成します。
• i を 0 から nums のサイズまで繰り返します:
  • nums[i] と nums[0] を入れ替えます。
  • nums[0] が未訪問であれば util(0) を呼び出します。
  • nums[0] を visited に追加します。
  • nums[i] と nums[0] を元に戻します。
• 最後に res を返します。

それでは、実際の実装を見て理解を深めましょう。

実装例

from math import sqrt
class Solution:
   def solve(self, nums):
      self.res = 0
      def util(i):
         if i + 1 == len(nums):
            self.res += 1
            return
         visited = set()
         for j in range(i + 1, len(nums)):
            s = nums[i] + nums[j]
            if s not in visited and int(sqrt(s)) ** 2 == s:
               visited.add(s)
               nums[i + 1], nums[j] = nums[j], nums[i + 1]
               util(i + 1)
               nums[i + 1], nums[j] = nums[j], nums[i + 1]
      visited = set()
      for i in range(len(nums)):
         nums[i], nums[0] = nums[0], nums[i]
         if nums[0] not in visited:
            util(0)
         visited.add(nums[0])
         nums[i], nums[0] = nums[0], nums[i]
      return self.res
ob = Solution()
nums = [2, 9, 7]
print(ob.solve(nums))

入力

[2, 9, 7]

出力

2

ポイントのまとめ

このアルゴリズムの肝は、次の2点です。

• 枝刈りによる高速化: 同じ合計値を持つ候補は結果が重複するだけなので、visited 集合で一度だけ探索すれば十分です。
• 状態の復元: 再帰から戻った後に入れ替えを必ず元に戻すことで、他の経路の探索に影響を与えません。

この手法により、全順列を素朴に生成する場合(O(n!))に比べて、大幅に計算量を抑えながら正確に答えを求めることができます。