C++で解く平方数配列（Squareful配列）となる順列の個数

問題概要

正整数からなる配列 A が与えられたとします。すべての隣接する 2 要素の和が完全平方数であるとき、この配列を「平方数配列（Squareful 配列）」と呼びます。求めたいのは、A を並べ替えて作られる順列のうち、平方数配列となっているものの総数です。なお、あるインデックス i において A1[i] ≠ A2[i] となるとき、またそのときに限り、2 つの順列 A1 と A2 は「異なる」ものとみなします。

たとえば、入力が [3, 30, 6] のとき、答えは 2 になります。[3, 6, 30] と [30, 6, 3] という 2 通りの順列が条件を満たすためです（3 + 6 = 9、6 + 30 = 36 はいずれも完全平方数です）。

解き方のアプローチ

この問題は、バックトラッキング（引き戻し法）ですべての並べ替えを試しつつ、条件を満たさない枝を途中で切り捨てることで効率的に解けます。ポイントは次の 2 つです。

• 完全平方数の判定： 関数 isSqr(n) は n の平方根 x を求め、x * x == n が成立すれば true を返します。
• 重複順列の排除： 同じ位置に同じ値を置く操作は結果が重複するため、各再帰段階で set 型の visited を用意し、すでに試した値をスキップします。

手順の詳細

1. 関数 isSqr(n)： x := √n とし、(x * x) == n であれば true を返します。
2. 関数 solve(a, idx)：
   • idx が配列 a のサイズと等しい場合、count を 1 増やして return します。
   • 集合 visited を用意します。
   • i := idx から配列の末尾までループします。
     • (idx == 0 または isSqr(a[idx - 1] + a[i])) かつ a[i] が visited に存在しない場合：
       • a[idx] と a[i] を交換します。
       • solve(a, idx + 1) を再帰呼び出しします。
       • a[idx] と a[i] を元に戻します（バックトラック）。
       • a[i] を visited に追加します。
3. メイン処理： count := 0 で初期化し、solve(a, 0) を実行した後、count を返します。

最悪の場合の探索量は順列の総数に依存しますが、隣接要素の和による早期の枝刈りと重複値のスキップにより、実際に探索すべき状態は大幅に絞り込まれます。

C++ 実装例

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
   public:
   int count;
   bool isSqr(lli n){
      lli x = sqrt(n);
      return x * x == n;
   }
   void solve(vector<int>& a, int idx){
      if (idx == a.size()) {
         count++;
         return;
      }
      set<int> visited;
      for (int i = idx; i < a.size(); i++) {
         if ((idx == 0 || isSqr(a[idx - 1] + a[i])) &&
         !visited.count(a[i])) {
            swap(a[idx], a[i]);
            solve(a, idx + 1);
            swap(a[idx], a[i]);
            visited.insert(a[i]);
         }
      }
   }
   int numSquarefulPerms(vector<int>& a){
      count = 0;
      solve(a, 0);
      return count;
   }
};
main(){
   Solution ob;
   vector<int> v = {3,30,6};
   cout << (ob.numSquarefulPerms(v));
}

入力

{3,30,6}

出力

2