题解 | 正方形的数目

#include <bits/stdc++.h>
using namespace std;

int n;
int ans = 0;
vector<int>arr;

bool f(int a, int b){    //判断完全平方
    int s = a+b;     
    int r = sqrt(s);    //该步不可跳过
    if(r*r == s){      //不能写成sqrt(s)*sqrt(s) == s
        return true;
    }
    return false;
}

//搜索
void back(int num, int pos, vector<int>& visit){  
    if(num == n){      //选完所有的，能到这一步即是符合条件排列
        ans ++;
        return;
    }
    int flag = 0;   //用来判断是不是第一个选的
    int pre;   //保存上一个选的
    for(int i = 0; i < n; i++){
        if(visit[i]){     //筛选出选过的数
            continue;
        }
        if(flag == 0){    //第一个选的
            flag = 1;
        }
        else{    //之后选的
            if(arr[i] == arr[pre]){     //跳过重复的数
                continue;
            }
        }
        pre = i;   //保存上一个选的
        if(f(arr[pos], arr[i])){    //满足条件
            visit[i] = 1;
            back(num+1, i, visit);    //搜索下一个数
            visit[i] = 0;
        }
    }
}

int main() {
    cin >> n;
    arr.resize(n);
    for(int i = 0; i < n;i ++){
        cin >> arr[i];
    }
    sort(arr.begin(), arr.end());     //排序方便去重 
    vector<int>visit(n, 0);
    for(int i = 0; i < n; i++){
        if(i > 0 && arr[i] == arr[i-1]){   //去重
            continue;
        }
        visit[i] = 1;    
        back(1, i, visit);      //将arr[i]作为第一个
        visit[i] = 0;
    }
    cout << ans << endl;
    return 0;
}